https://github.com/google/jax
Tip revision: df2dffbc7c0732c43eadcbeb3663b2febbb2037e authored by jax authors on 08 July 2024, 17:24:40 UTC
Sort devices explicitly by process index, then id (as opposed to IDs alone). IDs may be randomly generated, and are not guaranteed to be ordered based on their process index.
Sort devices explicitly by process index, then id (as opposed to IDs alone). IDs may be randomly generated, and are not guaranteed to be ordered based on their process index.
Tip revision: df2dffb
| File | Mode | Size |
|---|---|---|
| _static | ||
| _templates | ||
| _tutorials | ||
| cuda_custom_call | ||
| debugging | ||
| export | ||
| 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-- | 10.2 KB |
| api_compatibility.md | -rw-r--r-- | 3.0 KB |
| async_dispatch.rst | -rw-r--r-- | 5.3 KB |
| autodidax.ipynb | -rw-r--r-- | 140.7 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.3 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.1 KB |
| developer.md | -rw-r--r-- | 28.7 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.5 KB |
| index.rst | -rw-r--r-- | 2.3 KB |
| installation.md | -rw-r--r-- | 12.0 KB |
| investigating_a_regression.md | -rw-r--r-- | 6.6 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.pallas.rst | -rw-r--r-- | 292 bytes |
| jax.experimental.pjit.rst | -rw-r--r-- | 138 bytes |
| jax.experimental.rst | -rw-r--r-- | 928 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.export.rst | -rw-r--r-- | 1.0 KB |
| jax.extend.ffi.rst | -rw-r--r-- | 166 bytes |
| 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-- | 230 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-- | 488 bytes |
| jax.nn.initializers.rst | -rw-r--r-- | 841 bytes |
| jax.nn.rst | -rw-r--r-- | 675 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.4 KB |
| jax.scipy.rst | -rw-r--r-- | 6.6 KB |
| jax.sharding.rst | -rw-r--r-- | 572 bytes |
| jax.stages.rst | -rw-r--r-- | 448 bytes |
| jax.tree.rst | -rw-r--r-- | 261 bytes |
| jax.tree_util.rst | -rw-r--r-- | 778 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.0 KB |
| jit-compilation.md | -rw-r--r-- | 11.0 KB |
| key-concepts.md | -rw-r--r-- | 6.1 KB |
| multi_process.md | -rw-r--r-- | 8.5 KB |
| notes.rst | -rw-r--r-- | 1.1 KB |
| persistent_compilation_cache.md | -rw-r--r-- | 8.3 KB |
| profiling.md | -rw-r--r-- | 12.5 KB |
| pytrees.md | -rw-r--r-- | 10.8 KB |
| quickstart.md | -rw-r--r-- | 8.4 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.5 KB |
| sharded-computation.md | -rw-r--r-- | 12.1 KB |
| stateful-computations.md | -rw-r--r-- | 7.9 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-- | 712 bytes |
| working-with-pytrees.md | -rw-r--r-- | 17.2 KB |
