Merge pull request #22794 from jakevdp:array-api-astype

PiperOrigin-RevId: 658456484
jax-ml · Aug 1, 2024 · b3924da · b3924da
2 parents 06ffa70 + c24f209
commit b3924da
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 41 deletions.
diff --git a/jax/experimental/array_api/__init__.py b/jax/experimental/array_api/__init__.py
@@ -58,6 +58,7 @@
     asarray as asarray,
     asin as asin,
     asinh as asinh,
+    astype as astype,
     atan as atan,
     atan2 as atan2,
     atanh as atanh,
@@ -190,7 +191,3 @@
     zeros as zeros,
     zeros_like as zeros_like,
 )
-
-from jax.experimental.array_api._data_type_functions import (
-    astype as astype,
-)
diff --git a/jax/experimental/array_api/_data_type_functions.py b/jax/experimental/array_api/_data_type_functions.py