https://github.com/google/jax
Tip revision: bbfd7e044a5f6de96d3a0bb60d74be764bb48513 authored by Zhuo Peng on 12 August 2024, 18:01:51 UTC
Fixed a bug that may occur with the backward pass of a `tf_fun` wrapped by `jax2tf.call_tf`, when `tf_fun` mutates the structure of its input parameters.
Fixed a bug that may occur with the backward pass of a `tf_fun` wrapped by `jax2tf.call_tf`, when `tf_fun` mutates the structure of its input parameters.
Tip revision: bbfd7e0
File | Mode | Size |
---|---|---|
_static | ||
_templates | ||
_tutorials | ||
cuda_custom_call | ||
debugging | ||
export | ||
ffi | ||
gpu_ops | ||
jep | ||
notebooks | ||
pallas | ||
sphinxext | ||
Custom_Operation_for_GPUs.md | -rw-r--r-- | 42.7 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-- | 1013 bytes |
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.4 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.2 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 |
ffi.ipynb | -rw-r--r-- | 30.3 KB |
ffi.md | -rw-r--r-- | 24.0 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-- | 889 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.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-- | 440 bytes |
jax.experimental.pjit.rst | -rw-r--r-- | 138 bytes |
jax.experimental.rst | -rw-r--r-- | 902 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-- | 188 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-- | 10.0 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-- | 801 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-- | 526 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-- | 719 bytes |
working-with-pytrees.md | -rw-r--r-- | 17.2 KB |