How to identify the batch dimension?

[mattiasmar]

How can I in a programmatical way tell what dimension of a tensor represents the batch dimension?

For example the linen Conv op, when defined like this:
nn.Conv(features=32, kernel_size=(3, 3))
generates a bias tensor of shape (32,) and a kernel of shape (3, 3, 1, 32).
What (or where) is the logic that causes the batch to be assigned the 3rd dimension?

[soraros]

Maybe the Flax repo is the better place to ask about this.
But in this particular case(nn.Conv), you might want to check out this comment.