JAX compared to PyTorch 2: Get a feeling for JAX!

A simple torch.nn.Module for neural network model definition and training with gradient descent in PyTorch2 compared to a similar code implementation in JAX, in functional programming.
How to convert a stateful to a stateless operation in JAX, in functional programming? A simple coding example in JAX: regression via gradient descent, where there is one kind of state: the model parameters.
Link to documentation and free Colab NB:
jax.readthedoc...
colab.research...
#jax
#ai
#parallel
#computerscience
#computertipsandtricks

[@MaxwellHay]

I like the way jax handles the dependencies explicitly from a software engineer point of view.

[@EffectCrash]

love every video you made!

[@riser9644]

Very interesting video

[@riser9644]

Beyond experiment and curiosity, is there any benifit to JAX and it's variebts Trax,etc over pytorch, Tf