Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ENH: is_lazy_array() #228

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Conversation

crusaderky
Copy link
Contributor

Closes #225

@@ -92,6 +93,70 @@ def test_is_writeable_array_numpy():
assert not is_writeable_array(x)


@pytest.mark.parametrize("library", all_libraries)
Copy link
Contributor Author

@crusaderky crusaderky Jan 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#227 / #232 increase coverage

Comment on lines +855 to +862
# Unknown Array API compatible object. Note that this test may have dire consequences
# in terms of performance, e.g. for a lazy object that eagerly computes the graph
# on __bool__ (dask is one such example, which however is special-cased above).
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a good argument for replacing everything below this point with a blind return False. Please discuss if you'd prefer it that way.

if (
is_numpy_array(x)
or is_cupy_array(x)
or is_torch_array(x)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FIXME is this correct and safe? I don't know enough about torch.compile myself.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It works, at least for most backends, since TorchDynamo can handle graph breaks. If we have to make a choice here, then I'd say what you wrote is correct for PyTorch.

@crusaderky crusaderky mentioned this pull request Jan 7, 2025
@crusaderky crusaderky closed this Jan 7, 2025
@crusaderky crusaderky reopened this Jan 7, 2025
tests/test_common.py Outdated Show resolved Hide resolved
# Unknown Array API compatible object. Note that this test may have dire consequences
# in terms of performance, e.g. for a lazy object that eagerly computes the graph
# on __bool__ (dask is one such example, which however is special-cased above).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

array_api_strict reaches this point

tests/test_common.py Outdated Show resolved Hide resolved
tests/test_common.py Outdated Show resolved Hide resolved
Comment on lines +105 to +130
@pytest.mark.parametrize("shape", [(math.nan,), (1, math.nan), (None, ), (1, None)])
def test_is_lazy_array_nan_size(shape, monkeypatch):
"""Test is_lazy_array() on an unknown Array API compliant object
with NaN (like Dask) or None (like ndonnx) in its shape
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be simplified after #231


# Select a single point of the array
s = size(x)
if s is None or math.isnan(s):
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To be simplified after #231

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Conditionally run health checks on jitted JAX arrays and dask arrays
3 participants