mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 12:26:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #18338 from froystig:partitionable-threefry-ctx-mgr

Browse Source 
PiperOrigin-RevId: 578325021
This commit is contained in:
jax authors 2023-10-31 15:44:35 -07:00
commit 49fedb1c52
2 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/__init__.py
View File

@ -67,6 +67,7 @@ from jax._src.config import (
numpy_rank_promotion as numpy_rank_promotion,
jax2tf_associative_scan_reductions as jax2tf_associative_scan_reductions,
legacy_prng_key as legacy_prng_key,
threefry_partitionable as threefry_partitionable,
transfer_guard as transfer_guard,
transfer_guard_host_to_device as transfer_guard_host_to_device,
transfer_guard_device_to_device as transfer_guard_device_to_device,

1
jax/_src/config.py
View File

@ -683,6 +683,7 @@ class _ThreadLocalExtraJitContext(NamedTuple):
numpy_dtype_promotion: Optional[str] = None
default_matmul_precision: Optional[Any] = None
dynamic_shapes: bool = False
threefry_partitionable: bool = False
softmax_custom_jvp: bool = False
xla_profile_version: int = 0