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

DOC: make clear that printoptions are NumPy aliases

Browse Source
This commit is contained in:
Jake VanderPlas 2025-02-07 09:56:52 -08:00
parent c521bc6205
commit 08563842b9
2 changed files with 35 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -111,11 +111,40 @@ euler_gamma = np.euler_gamma
inf = np.inf
nan = np.nan
# NumPy utility functions
# Wrappers for NumPy printoptions
get_printoptions = np.get_printoptions
printoptions = np.printoptions
set_printoptions = np.set_printoptions
def get_printoptions():
"""Alias of :func:`numpy.get_printoptions`.
JAX arrays are printed via NumPy, so NumPy's `printoptions`
configurations will apply to printed JAX arrays.
See the :func:`numpy.set_printoptions` documentation for details
on the available options and their meanings.
"""
return np.get_printoptions()
def printoptions(*args, **kwargs):
"""Alias of :func:`numpy.printoptions`.
JAX arrays are printed via NumPy, so NumPy's `printoptions`
configurations will apply to printed JAX arrays.
See the :func:`numpy.set_printoptions` documentation for details
on the available options and their meanings.
"""
return np.printoptions(*args, **kwargs)
def set_printoptions(*args, **kwargs):
"""Alias of :func:`numpy.set_printoptions`.
JAX arrays are printed via NumPy, so NumPy's `printoptions`
configurations will apply to printed JAX arrays.
See the :func:`numpy.set_printoptions` documentation for details
on the available options and their meanings.
"""
return np.set_printoptions(*args, **kwargs)
@export
def iscomplexobj(x: Any) -> bool:

3
tests/lax_numpy_test.py
View File

@ -6567,7 +6567,8 @@ class NumpyDocTests(jtu.JaxTestCase):
aliases = ['abs', 'acos', 'acosh', 'asin', 'asinh', 'atan', 'atanh', 'atan2',
'amax', 'amin', 'around', 'bitwise_invert', 'bitwise_left_shift',
'bitwise_not','bitwise_right_shift', 'conj', 'degrees', 'divide',
'mod', 'pow', 'radians', 'round_']
'get_printoptions', 'mod', 'pow', 'printoptions', 'radians', 'round_',
'set_printoptions']
skip_args_check = ['vsplit', 'hsplit', 'dsplit', 'array_split']
for name in dir(jnp):