rockpool.nn.modules.jax
=======================

.. automodule:: rockpool.nn.modules.jax

   
   
   

   
   
   

   
   
   .. rubric:: Classes

   .. autosummary::
   
      Array
   
   

   
   
   



.. rubric:: Modules

.. autosummary::
   :toctree:
   :recursive:

   rockpool.nn.modules.jax.abstract_arrays
   rockpool.nn.modules.jax.ad_checkpoint
   rockpool.nn.modules.jax.api_util
   rockpool.nn.modules.jax.cloud_tpu_init
   rockpool.nn.modules.jax.collect_profile
   rockpool.nn.modules.jax.config
   rockpool.nn.modules.jax.core
   rockpool.nn.modules.jax.custom_batching
   rockpool.nn.modules.jax.custom_derivatives
   rockpool.nn.modules.jax.custom_transpose
   rockpool.nn.modules.jax.debug
   rockpool.nn.modules.jax.distributed
   rockpool.nn.modules.jax.dlpack
   rockpool.nn.modules.jax.dtypes
   rockpool.nn.modules.jax.errors
   rockpool.nn.modules.jax.example_libraries
   rockpool.nn.modules.jax.experimental
   rockpool.nn.modules.jax.flatten_util
   rockpool.nn.modules.jax.image
   rockpool.nn.modules.jax.interpreters
   rockpool.nn.modules.jax.jaxpr_util
   rockpool.nn.modules.jax.lax
   rockpool.nn.modules.jax.lib
   rockpool.nn.modules.jax.linear_util
   rockpool.nn.modules.jax.monitoring
   rockpool.nn.modules.jax.nn
   rockpool.nn.modules.jax.numpy
   rockpool.nn.modules.jax.ops
   rockpool.nn.modules.jax.prng
   rockpool.nn.modules.jax.profiler
   rockpool.nn.modules.jax.random
   rockpool.nn.modules.jax.scipy
   rockpool.nn.modules.jax.sharding
   rockpool.nn.modules.jax.stages
   rockpool.nn.modules.jax.test_util
   rockpool.nn.modules.jax.tools
   rockpool.nn.modules.jax.tree_util
   rockpool.nn.modules.jax.typing
   rockpool.nn.modules.jax.util
   rockpool.nn.modules.jax.version