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

DOC: add caveats to jnp.ndarray.astype

Browse Source
This commit is contained in:
Jake VanderPlas 2022-03-23 09:38:15 -07:00
parent e0d3946e34
commit 852a747189
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -4341,6 +4341,13 @@ def nanmedian(a, axis: Optional[Union[int, Tuple[int, ...]]] = None, out=None,
def _astype(arr, dtype):
"""Copy the array and cast to a specified dtype.
This is implemeted via :func:`jax.lax.convert_element_type`, which may
have slightly different behavior than :meth:`numpy.ndarray.astype` in
some cases. In particular, the details of float-to-int and int-to-float
casts are implementation dependent.
"""
if dtype is None:
dtype = dtypes.canonicalize_dtype(float_)
lax_internal._check_user_dtype_supported(dtype, "astype")