Gemfury

duality-group / jax python

Repository URL to install this package:

Details

jax / jax / _src
..
clusters
debugger
image
internal_test_util
interpreters
lax
lib
nn
numpy
ops
scipy
state
third_party
__init__.py
abstract_arrays.py
ad_checkpoint.py
ad_util.py
api.py
api_util.py
array.py
basearray.py
basearray.pyi
callback.py
checkify.py
cloud_tpu_init.py
compilation_cache.py
compilation_cache_interface.py
config.py
core.py
custom_api_util.py
custom_batching.py
custom_derivatives.py
custom_transpose.py
debugging.py
deprecations.py
dispatch.py
distributed.py
dlpack.py
dtypes.py
effects.py
environment_info.py
errors.py
flatten_util.py
gfile_cache.py
lax_reference.py
lazy_loader.py
linear_util.py
maps.py
mesh.py
monitoring.py
op_shardings.py
partition_spec.py
path.py
pickle_util.py
pjit.py
pretty_printer.py
prng.py
profiler.py
public_test_util.py
random.py
sharding.py
sharding_impls.py
sharding_specs.py
source_info_util.py
stages.py
test_util.py
tpu_custom_call.py
traceback_util.py
tree_util.py
typing.py
util.py
xla_bridge.py