mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

add jax.threefry_partitionable context manager

Browse Source
This commit is contained in:
Roy Frostig 2023-10-31 13:42:45 -07:00
parent b22e75716f
commit ed9a4c2939
1 changed files with 1 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,