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