deleted

Btw, regarding Jax, I think of Mathematica implemented tensor differentiation natively, this would obviate the need to rely on external autodiff.

The issue is that neural network training involves differentiation of expression like norm(W1 W2 W3) w.r.t to matrix W2. Mathematica differentiation is at its core a scalar differentiation, it can't treat W2 as an atom. Autodiff frameworks like Jax can find these derivatives efficiently by treating matmuls as atomic operations.

I have thought a lot about the proper design of such a system and have a couple of prototypes, getting it into Mathematica is on my to-do list https://community.wolfram.com/groups/-/m/t/2437093

Just wanted to alert that as of mid 2023 Google is recommending Flax for new projects instead of Haiku

What machine learning framework (or frameworks) are of interest to you (as a possible complete replacement for MXNet)?

Building on Yaroslav Bulatov's post:

Closely related is support for things like NeuralODEs (backpropogating through NDSolve). This is another increasingly important area in scientific ML, and is now built into Julia (https://julialang.org/blog/2019/01/fluxdiffeq/) and PyTorch (https://towardsdatascience.com/neural-odes-with-pytorch-lightning-and-torchdyn-87ca4a7c6ffd) . One can kind of "hack it" by implementing ODE integrators as neural nets in Mathematica (I've done it, as have some others).