Getting Started
Reference Documentation
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
jit
vmap
pmap
jax.nn.
glu
Gated linear unit activation function.
x (Any) – input array
Any
axis (int) – the axis along which the split should be computed (default: -1)
int