https://github.com/google/jax
Tip revision: b06663d0d366bf158ebde3fce4b55e8ee7312a07 authored by jax authors on 22 May 2024, 14:57:25 UTC
Merge pull request #21358 from hawkinsp:types
Merge pull request #21358 from hawkinsp:types
Tip revision: b06663d
| File | Mode | Size |
|---|---|---|
| _static | ||
| _templates | ||
| _tutorials | ||
| debugging | ||
| gpu_ops | ||
| jep | ||
| notebooks | ||
| pallas | ||
| sphinxext | ||
| Custom_Operation_for_GPUs.md | -rw-r--r-- | 42.8 KB |
| Custom_Operation_for_GPUs.py | -rw-r--r-- | 18.8 KB |
| README.md | -rw-r--r-- | 133 bytes |
| advanced_guide.rst | -rw-r--r-- | 1.0 KB |
| aot.md | -rw-r--r-- | 9.3 KB |
| api_compatibility.md | -rw-r--r-- | 2.9 KB |
| async_dispatch.rst | -rw-r--r-- | 5.3 KB |
| autodidax.ipynb | -rw-r--r-- | 140.6 KB |
| autodidax.md | -rw-r--r-- | 103.1 KB |
| autodidax.py | -rw-r--r-- | 103.0 KB |
| automatic-differentiation.md | -rw-r--r-- | 8.3 KB |
| automatic-vectorization.md | -rw-r--r-- | 3.6 KB |
| beginner_guide.rst | -rw-r--r-- | 2.0 KB |
| build_custom_gpu.sh | -rw-r--r-- | 1.3 KB |
| building_on_jax.md | -rw-r--r-- | 4.6 KB |
| changelog.md | -rw-r--r-- | 33 bytes |
| concurrency.rst | -rw-r--r-- | 609 bytes |
| conf.py | -rw-r--r-- | 11.1 KB |
| contributing.md | -rw-r--r-- | 7.0 KB |
| contributor_guide.rst | -rw-r--r-- | 395 bytes |
| debugging.md | -rw-r--r-- | 5.8 KB |
| deprecation.md | -rw-r--r-- | 2.0 KB |
| developer.md | -rw-r--r-- | 27.0 KB |
| device_memory_profiling.md | -rw-r--r-- | 4.8 KB |
| distributed_data_loading.md | -rw-r--r-- | 25.4 KB |
| errors.rst | -rw-r--r-- | 480 bytes |
| faq.rst | -rw-r--r-- | 36.5 KB |
| glossary.rst | -rw-r--r-- | 4.5 KB |
| gpu_memory_allocation.rst | -rw-r--r-- | 2.8 KB |
| gpu_performance_tips.md | -rw-r--r-- | 5.2 KB |
| index.rst | -rw-r--r-- | 2.3 KB |
| installation.md | -rw-r--r-- | 12.4 KB |
| investigating_a_regression.md | -rw-r--r-- | 6.5 KB |
| jax.debug.rst | -rw-r--r-- | 630 bytes |
| jax.distributed.rst | -rw-r--r-- | 196 bytes |
| jax.dlpack.rst | -rw-r--r-- | 178 bytes |
| jax.dtypes.rst | -rw-r--r-- | 226 bytes |
| jax.example_libraries.optimizers.rst | -rw-r--r-- | 195 bytes |
| jax.example_libraries.rst | -rw-r--r-- | 943 bytes |
| jax.example_libraries.stax.rst | -rw-r--r-- | 177 bytes |
| jax.experimental.array_api.rst | -rw-r--r-- | 120 bytes |
| jax.experimental.checkify.rst | -rw-r--r-- | 345 bytes |
| jax.experimental.compilation_cache.rst | -rw-r--r-- | 333 bytes |
| jax.experimental.custom_partitioning.rst | -rw-r--r-- | 198 bytes |
| jax.experimental.host_callback.rst | -rw-r--r-- | 257 bytes |
| jax.experimental.jet.rst | -rw-r--r-- | 134 bytes |
| jax.experimental.key_reuse.rst | -rw-r--r-- | 120 bytes |
| jax.experimental.maps.rst | -rw-r--r-- | 167 bytes |
| jax.experimental.mesh_utils.rst | -rw-r--r-- | 225 bytes |
| jax.experimental.multihost_utils.rst | -rw-r--r-- | 453 bytes |
| jax.experimental.pjit.rst | -rw-r--r-- | 138 bytes |
| jax.experimental.rst | -rw-r--r-- | 1015 bytes |
| jax.experimental.serialize_executable.rst | -rw-r--r-- | 241 bytes |
| jax.experimental.shard_map.rst | -rw-r--r-- | 185 bytes |
| jax.experimental.sparse.rst | -rw-r--r-- | 2.3 KB |
| jax.extend.linear_util.rst | -rw-r--r-- | 264 bytes |
| jax.extend.mlir.rst | -rw-r--r-- | 163 bytes |
| jax.extend.random.rst | -rw-r--r-- | 266 bytes |
| jax.extend.rst | -rw-r--r-- | 213 bytes |
| jax.flatten_util.rst | -rw-r--r-- | 224 bytes |
| jax.image.rst | -rw-r--r-- | 332 bytes |
| jax.lax.rst | -rw-r--r-- | 4.2 KB |
| jax.lib.rst | -rw-r--r-- | 458 bytes |
| jax.nn.initializers.rst | -rw-r--r-- | 841 bytes |
| jax.nn.rst | -rw-r--r-- | 630 bytes |
| jax.numpy.rst | -rw-r--r-- | 8.9 KB |
| jax.ops.rst | -rw-r--r-- | 489 bytes |
| jax.profiler.rst | -rw-r--r-- | 656 bytes |
| jax.random.rst | -rw-r--r-- | 1.0 KB |
| jax.rst | -rw-r--r-- | 2.3 KB |
| jax.scipy.rst | -rw-r--r-- | 6.6 KB |
| jax.sharding.rst | -rw-r--r-- | 644 bytes |
| jax.stages.rst | -rw-r--r-- | 441 bytes |
| jax.tree.rst | -rw-r--r-- | 261 bytes |
| jax.tree_util.rst | -rw-r--r-- | 659 bytes |
| jax.typing.rst | -rw-r--r-- | 178 bytes |
| jax_array_migration.md | -rw-r--r-- | 11.8 KB |
| jax_internal_api.rst | -rw-r--r-- | 161 bytes |
| jaxpr.rst | -rw-r--r-- | 17.3 KB |
| jit-compilation.md | -rw-r--r-- | 10.9 KB |
| key-concepts.md | -rw-r--r-- | 6.1 KB |
| multi_process.md | -rw-r--r-- | 9.1 KB |
| notes.rst | -rw-r--r-- | 1.1 KB |
| persistent_compilation_cache.md | -rw-r--r-- | 2.5 KB |
| profiling.md | -rw-r--r-- | 12.4 KB |
| pytrees.md | -rw-r--r-- | 10.7 KB |
| quickstart.md | -rw-r--r-- | 8.5 KB |
| random-numbers.md | -rw-r--r-- | 8.2 KB |
| rank_promotion_warning.rst | -rw-r--r-- | 1.8 KB |
| requirements.txt | -rw-r--r-- | 491 bytes |
| sharded-computation.ipynb | -rw-r--r-- | 39.4 KB |
| sharded-computation.md | -rw-r--r-- | 12.0 KB |
| stateful-computations.md | -rw-r--r-- | 7.8 KB |
| transfer_guard.rst | -rw-r--r-- | 2.7 KB |
| tutorials.rst | -rw-r--r-- | 287 bytes |
| type_promotion.rst | -rw-r--r-- | 14.0 KB |
| user_guides.rst | -rw-r--r-- | 696 bytes |
| working-with-pytrees.md | -rw-r--r-- | 17.2 KB |
