Seeded Mini-Batch Indices

Problem

Build the index batches for one epoch of mini-batch SGD.

Start from indices [0, 1, ..., n-1], shuffle them after random.seed(seed) plus random.shuffle, then split them into consecutive chunks of batch_size. The last chunk may be shorter, and an empty input yields []. Return a list of index lists.

Hint

Reseed before shuffling so the batch order is reproducible for a given seed.

Test cases

case 1

in: [6,2,1]

expected: [[2,3],[5,0],[4,1]]

case 2

in: [7,3,42]

expected: [[1,3,4],[2,6,0],[5]]

case 3

in: [3,5,0]

expected: [[0,2,1]]

case 4

in: [0,4,7]

expected: []

Starter codePython

def mini_batch_indices(n, batch_size, seed):
    # Your code here
    pass

4 test cases · run your code against them in the editor