Module phi.jax.stax
Sub-modules
phi.jax.stax.flow-
Standard import for Jax + Stax mode …
phi.jax.stax.nets-
Stax implementation of the unified machine learning API. Equivalent functions also exist for the other frameworks …
phi.jax.staxphi.jax.stax.flowStandard import for Jax + Stax mode …
phi.jax.stax.netsStax implementation of the unified machine learning API. Equivalent functions also exist for the other frameworks …