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

[array API] make capabilities more accurate

Browse Source
This commit is contained in:
Jake VanderPlas 2025-03-26 12:11:47 -07:00
parent ec2f0f5913
commit 096810a721
2 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jax/_src/numpy/array_api_metadata.py
View File

@ -51,8 +51,9 @@ class ArrayNamespaceInfo:
.. _Python array API: https://data-apis.org/array-api/
"""
_capabilities = {
"boolean indexing": True,
"data-dependent shapes": False,
"boolean indexing": False, # within transformations
"data-dependent shapes": False, # within transformations
"max dimensions": 64, # XLA limitation
}
def _build_dtype_dict(self):

3
tests/array_api_test.py
View File

@ -275,8 +275,9 @@ class ArrayAPIInspectionUtilsTest(jtu.JaxTestCase):
def test_capabilities_info(self):
capabilities = self.info.capabilities()
assert capabilities["boolean indexing"]
assert not capabilities["boolean indexing"]
assert not capabilities["data-dependent shapes"]
assert capabilities["max dimensions"] == 64
def test_default_device_info(self):
assert self.info.default_device() is None