TBH, I don't understand the jax code, it must rely on some library I don't know, stateless_call, attach_weights, etc don't seem very functional to me. In general, my implementation would look simpler: just passing in the weights from outside and take derivatives.
-
This Tweet is unavailable.
-
If I don't have any non-trainable weights, I get this, which is much simpler. However, adding non-trainable weights in the pictures adds significant complication. The same applies to metric state updates as well.pic.twitter.com/oQ8zbfoq3W
Maybe there are better solutions, I don't know. I'm still learning JAX.
-
-
The Linen version hides everything in the "optimizer", even the weights and gives back a new "optimizer". One might say it is a misnomer. Then one can replace various parts of the optimizer and create new optimizers.
- Show replies
New conversation -
Loading seems to be taking a while.
Twitter may be over capacity or experiencing a momentary hiccup. Try again or visit Twitter Status for more information.