README.md
setup.py
examples/__init__.py
examples/datasets.py
examples/examples_test.py
examples/kernel_lsq.py
examples/mnist_classifier.py
examples/mnist_classifier_fromscratch.py
examples/mnist_vae.py
examples/onnx2xla.py
examples/prng_compile_time.py
examples/resnet50.py
jax/__init__.py
jax/abstract_arrays.py
jax/ad_util.py
jax/api.py
jax/api_util.py
jax/config.py
jax/core.py
jax/lax.py
jax/lax_linalg.py
jax/lax_reference.py
jax/linear_util.py
jax/pprint_util.py
jax/random.py
jax/test_util.py
jax/tree_util.py
jax/util.py
jax.egg-info/PKG-INFO
jax.egg-info/SOURCES.txt
jax.egg-info/dependency_links.txt
jax.egg-info/requires.txt
jax.egg-info/top_level.txt
jax/experimental/__init__.py
jax/experimental/lapax.py
jax/experimental/minmax.py
jax/experimental/stax.py
jax/interpreters/__init__.py
jax/interpreters/ad.py
jax/interpreters/batching.py
jax/interpreters/partial_eval.py
jax/interpreters/xla.py
jax/lib/__init__.py
jax/lib/pywrap_xla.py
jax/lib/xla_bridge.py
jax/lib/xla_client.py
jax/lib/xla_data_pb2.py
jax/numpy/__init__.py
jax/numpy/fft.py
jax/numpy/lax_numpy.py
jax/numpy/linalg.py
jax/scipy/__init__.py
jax/scipy/linalg.py
jax/scipy/misc.py
jax/scipy/special.py
jax/scipy/stats/__init__.py
jax/scipy/stats/beta.py
jax/scipy/stats/expon.py
jax/scipy/stats/gamma.py
jax/scipy/stats/laplace.py
jax/scipy/stats/norm.py
jax/scipy/stats/uniform.py