The JAX library offers most things that you need for making neural networks, but there is no shortage of frameworks/libraries that build on JAX to cat

Search code, repositories, users, issues, pull requests...

submited by

Style Pass

2024-10-27 06:30:03

The JAX library offers most things that you need for making neural networks, but there is no shortage of frameworks/libraries that build on JAX to cater to neural net building.

A neural network $f$ is simply mathematical function of data $X$, parameters $\theta$, and hyper-parameters $\alpha$. We place $\theta$ as the first parameter of $f$ because jax.grad creates the gradient of $f$ wrt to the first parameter by default.

Notice that we named params whenever it was passed to the encoder mlp: params["mlp_1] and params["mlp_2"]. These names are essential and is part of zephyr's design to allow maximum control over all parameters.

Notice that an mlp is not some object, not some function passed to a transform, not a dataclass PyTree object, it is simply a function mlp(params, x, num_out_per_layer). There is no need to instatiate a model or a neural network. It's just a function! (Later we will show more reasons why this would be advantageous)

We have an autoencoder, now how do we instatiate the model? As said before, no instatiation needed. What we do need is a an initial params. This is easy with the trace function.