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

Expose PrecisionLike

Browse Source 
This is used in client code like:
https://github.com/search?q=repo%3Agoogle%2Fflax%20%20PrecisionLike&type=code
This commit is contained in:
Neil Girdhar 2023-12-05 10:50:22 -05:00
parent bcf3255fe6
commit 9f85beb56b
2 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/jax.lax.rst
View File

@ -252,5 +252,6 @@ Argument classes
.. autoclass:: GatherDimensionNumbers
.. autoclass:: GatherScatterMode
.. autoclass:: Precision
.. autoclass:: PrecisionLike
.. autoclass:: RoundingMethod
.. autoclass:: ScatterDimensionNumbers

1
jax/lax/__init__.py
View File

@ -18,6 +18,7 @@
from jax._src.lax.lax import (
DotDimensionNumbers as DotDimensionNumbers,
Precision as Precision,
PrecisionLike as PrecisionLike,
RandomAlgorithm as RandomAlgorithm,
RoundingMethod as RoundingMethod,
abs as abs,