I've deleted my tweet comparing different approaches to writing DL training loops because some folks told me it mad JAX look more difficult than it is. I'm still learning JAX (which has been a lot of fun), and if the JAX code I write is bad, that's entirely on me.
-
I was following the thread, it didn't seem too mad.. what i've learned in jax they expose the params and state as trainable and non-trainable parameters.. which is a pytree with values from the model and optimizer.
I don't think the code was wrong, but it could be made a bit simpler and more idiomatic. I need to learn to write idiomatic JAX code. Idiomatax code?
