Open
Description
All over Mava we have statements like tree.map(lambda x: x[jnp.newaxis], pytree) I think it would make it a lot clearer if we had an add_batch_dim method in utils.jax_utils
{{ message }}
No one assigned
No type
No projects
No milestone
None yet
No branches or pull requests