Python Module Index

f | j
 
f
flax
    flax.configurations
    flax.core.variables
    flax.errors
    flax.jax_utils
    flax.linen
    flax.linen.activation
    flax.linen.initializers
    flax.linen.spmd
    flax.linen.transforms
    flax.serialization
    flax.struct
    flax.traceback_util
    flax.training.checkpoints
    flax.training.lr_schedule
    flax.traverse_util
 
j
jax
    jax.core
    jax.debug
    jax.distributed
    jax.dlpack
    jax.dtypes
    jax.example_libraries
    jax.example_libraries.optimizers
    jax.example_libraries.stax
    jax.experimental.checkify
    jax.experimental.custom_partitioning
    jax.experimental.host_callback
    jax.experimental.jet
    jax.experimental.maps
    jax.experimental.multihost_utils
    jax.experimental.pjit
    jax.experimental.sparse
    jax.experimental.sparse.linalg
    jax.extend
    jax.extend.linear_util
    jax.extend.random
    jax.flatten_util
    jax.image
    jax.lax
    jax.lax.linalg
    jax.nn
    jax.nn.initializers
    jax.numpy
    jax.numpy.fft
    jax.numpy.linalg
    jax.ops
    jax.profiler
    jax.random
    jax.scipy.fft
    jax.scipy.integrate
    jax.scipy.linalg
    jax.scipy.ndimage
    jax.scipy.optimize
    jax.scipy.signal
    jax.scipy.sparse.linalg
    jax.scipy.spatial.transform
    jax.scipy.special
    jax.scipy.stats
    jax.scipy.stats.bernoulli
    jax.scipy.stats.beta
    jax.scipy.stats.betabinom
    jax.scipy.stats.binom
    jax.scipy.stats.cauchy
    jax.scipy.stats.chi2
    jax.scipy.stats.dirichlet
    jax.scipy.stats.expon
    jax.scipy.stats.gamma
    jax.scipy.stats.gennorm
    jax.scipy.stats.geom
    jax.scipy.stats.laplace
    jax.scipy.stats.logistic
    jax.scipy.stats.multinomial
    jax.scipy.stats.multivariate_normal
    jax.scipy.stats.nbinom
    jax.scipy.stats.norm
    jax.scipy.stats.pareto
    jax.scipy.stats.poisson
    jax.scipy.stats.t
    jax.scipy.stats.truncnorm
    jax.scipy.stats.uniform
    jax.scipy.stats.vonmises
    jax.scipy.stats.wrapcauchy
    jax.sharding
    jax.stages
    jax.tree_util
    jax.typing