Tutorials
Advanced JAX Tutorials
Notes
Developer documentation
API documentation
jit
vmap
pmap
jax.nn.
relu
Rectified linear unit activation function.
Computes the element-wise function:
x (Any) –
Any