MLP training example with xmap()

[che-shr-cat]

Hi!

I tried to create an MNIST classification example on MLP based on the Xmap tutorial (https://jax.readthedocs.io/en/latest/notebooks/xmap_tutorial.html).
My attempt is here:
https://colab.research.google.com/drive/1ZgskIH-vmNd73XZDjSgFV-_I6xbjEwpL?usp=sharing

The problem is that the tutorial does not contain anything with taking gradients of the code, and the code in the tutorial is incompatible with the grad() transformation because it relies on logsumexp, which uses lax.pmax, for which differentiation rule is not implemented.

A set of questions:

1. I understand that xmap() is still experimental, but what are its current status and future plans? Is xmap() mature enough to use it this way?
2. As my similar examples with pmap() and vmap() work correctly, I assume it should be pretty straightforward to use xmap() as well. Is it correct or I'm doing something fundamentally wrong, and xmap is not the right tool for it? Or maybe I'm missing some simple workaround here?
3. Can anyone suggest any working example of implementing MLP (or any other pure JAX neural network) with training using xmap()?

Thanks in advance!