Module phi.jax

Jax integration.

Importing this module registers the Jax backend with phiml.math. Without this, Jax tensors cannot be handled by phiml.math functions.

To make Jax the default backend, import phi.jax.flow.

Sub-modules

phi.jax.flow

Standard import for Jax mode without neural network support …

phi.jax.stax