How to batch data before training in Jax?

[bigyankarki]

I have a matrix of (4,31). eg [[31el], [31el], [31 array of 256 length each], [31el]]. Each of the four rows is a feature. And I have been trying to batch the data into say 10. I want to pad the unfilled batch with Zeros. However, none of the solutions I found in StackOverflow works. Is there any way to achieve this? Is there any way we can leverage Vamp for this?

[YouJiacheng]

IIUC, [[31el], [31el], [31 array of 256 length each], [31el]] is not a matrix of (4,31).
Is 31 is the batch size?

[jakevdp]

Thanks for the question – I think we'll need more detail about what your inputs are, and what the expected outputs are. You might refer to this page for tips on how to ask a question that is likely to get a good answer: https://stackoverflow.com/help/how-to-ask