I am trying to understand the implementation of NDArrayIter when using multiple inputs. In the documentation, the following example is given:

data = {‘data1’:np.zeros(shape=(10,2,2)), ‘data2’:np.zeros(shape=(20,2,2))}

label = {‘label1’:np.zeros(shape=(10,1)), ‘label2’:np.zeros(shape=(20,1))}

dataiter = mx.io.NDArrayIter(data, label, 3, True, last_batch_handle=‘discard’)

It seems that the batching dimension is always assumed to be there first, is this correct? How about if I wanted to batch some of the inputs but not others? If I wanted to do inference using the Scala interface, would I use NDArrayIter to build the DataBatch[s] and then use those to construct a custom DataBatch that included the non-batched input? I guess this could be done in a customer DataIter?