How to get per-example gradients in Flax?

[marcvanzee]

How to get per-example gradients in Flax?

[marcvanzee]

All layers accept batch dimension free inputs. Example:

def loss_fn(param, batch):
  return jnp.square(param * batch)

model = 3.
per_example_grad_fn = jax.vmap(jax.grad(loss_fn), in_axes=(None, 0))
per_example_grad_fn(param, jnp.ones((3,)))

[mu-arkhipov]

It would be nice to have some linen layers in the example like in this PR.

[VMS-6511]

Are there any examples of this with linen layers as mentioned above?