Merge pull request #17666 from jakevdp:jex-array-types

PiperOrigin-RevId: 566732103

jax/extend/__init__.py

@@ -29,6 +29,7 @@ Breaking changes will be announced via the
 """
 
 from jax.extend import (
+    core as core,
     linear_util as linear_util,
     random as random,
 )

jax/extend/core.py

@@ -0,0 +1,20 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/google/jax/issues/7570
+
+from jax._src.abstract_arrays import (
+  array_types as array_types
+)

tests/extend_test.py

@@ -16,6 +16,7 @@ from absl.testing import absltest
 
 import jax.extend as jex
 
+from jax._src import abstract_arrays
 from jax._src import linear_util
 from jax._src import prng
 from jax._src import test_util as jtu
@@ -36,6 +37,7 @@ class ExtendTest(jtu.JaxTestCase):
     self.assertIs(jex.random.unsafe_rbg_prng_impl, prng.unsafe_rbg_prng_impl)
 
     # Assume these are tested elsewhere, only check equivalence
+    self.assertIs(jex.core.array_types, abstract_arrays.array_types)
     self.assertIs(jex.linear_util.StoreException, linear_util.StoreException)
     self.assertIs(jex.linear_util.WrappedFun, linear_util.WrappedFun)
     self.assertIs(jex.linear_util.cache, linear_util.cache)