Skip to content

Commit ef88926

Browse files
committed
Wrap dask.array by default since it is not Array API Compatible
1 parent cdd1c8d commit ef88926

File tree

1 file changed

+5
-2
lines changed

1 file changed

+5
-2
lines changed

array_api_compat/common/_helpers.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -521,11 +521,14 @@ def your_function(x, y):
521521
import torch
522522
namespaces.add(torch)
523523
elif is_dask_array(x):
524-
if _use_compat:
524+
# dask main namespace is not array APi compatible
525+
# so return namespace from array-api-compat unless
526+
# explicitly requested otherwise
527+
if _use_compat or _use_compat is None:
525528
_check_api_version(api_version)
526529
from ..dask import array as dask_namespace
527530
namespaces.add(dask_namespace)
528-
else:
531+
elif _use_compat is False:
529532
import dask.array as da
530533
namespaces.add(da)
531534
elif is_jax_array(x):

0 commit comments

Comments
 (0)