Correct way to initialize model

[averageFlaxUser]

I initialize my model as follows:

phi_model.init(jax.random.PRNGKey(seed=2), jnp.ones([784]))['params']

My input is of dims=(100, 784) during training, there is anther input of dims=(300, 784). Essentially different batch sizes. What would be the correct method to initialize the mode?

... jnp.ones([784]))['params'] or ... jnp.ones([1, 784]))['params']?

[cgarciae]

Hey @averageFlaxUser, most Flax layers support multiple batch dimensions and even having no batch dimension. That said to be safe I usually initialize with a batch dimension of 1 so jnp.ones([1, 784]) is probably what I would use.

[averageFlaxUser]

Thank you. This makes sense.