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 precision warning and workaround to jnp.arange documentation

Browse Source
This commit is contained in:
alhridoy 2023-10-20 15:01:24 -06:00
parent d631fb10fd
commit 63f7cfe04c
1 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
jax/_src/numpy/lax_numpy.py
View File

@ -2317,7 +2317,16 @@ def identity(n: DimSize, dtype: DTypeLike | None = None) -> Array:
return eye(n, dtype=dtype)
@util._wraps(np.arange)
@util._wraps(np.arange,lax_description= """
.. note::
Using ``arange`` with the ``step`` argument can lead to precision errors,
especially with lower-precision data types like ``fp8`` and ``bf16``.
For more details, see the docstring of :func:`numpy.arange`.
To avoid precision errors, consider using an expression like
``(jnp.arange(-600, 600) * .01).astype(jnp.bfloat16)`` to generate a sequence in a higher precision
and then convert it to the desired lower precision.
""")
def arange(start: DimSize, stop: Optional[DimSize] = None,
step: DimSize | None = None, dtype: DTypeLike | None = None) -> Array:
dtypes.check_user_dtype_supported(dtype, "arange")