All modules for which code is available
- jax._src.lax.control_flow.loops
- jax._src.lax.lax
- jax._src.lax.slicing
- jax._src.nn.functions
- jax._src.numpy.lax_numpy
- jax._src.ops.special
- objax.constants
- objax.functional.core.activation
- objax.functional.core.ops
- objax.functional.core.pooling
- objax.functional.divergence
- objax.functional.loss
- objax.functional.parallel
- objax.gradient
- objax.io.checkpoint
- objax.io.ops
- objax.jaxboard
- objax.module
- objax.nn.init
- objax.nn.layers
- objax.optimizer.adam
- objax.optimizer.ema
- objax.optimizer.lars
- objax.optimizer.momentum
- objax.optimizer.sgd
- objax.privacy.dpsgd.gradient
- objax.privacy.dpsgd.privacyaccountant
- objax.random.random
- objax.util.image
- objax.util.objax2tf
- objax.util.tracing
- objax.util.util
- objax.variable
- objax.zoo.convnet
- objax.zoo.dnnet
- objax.zoo.resnet_v2
- objax.zoo.rnn
- objax.zoo.vgg
- objax.zoo.wide_resnet