Merge pull request #33423 from jakevdp:dlpack-imports

PiperOrigin-RevId: 834466701

Author: Google-ML-Automation

jax/dlpack.py

@@ -12,9 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-import jax._src.dlpack
-
 from jax._src.dlpack import (
   from_dlpack as from_dlpack,
   is_supported_dtype as is_supported_dtype,
@@ -33,5 +30,5 @@
 
 
 from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
-__getattr__ = jax._src.deprecations.deprecation_getattr(__name__, _deprecations)
+__getattr__ = _deprecation_getattr(__name__, _deprecations)
 del _deprecation_getattr

tests/documentation_coverage_test.py

@@ -57,7 +57,6 @@ def jax_docs_dir() -> str:
   'jax.custom_transpose': ['custom_transpose'],
   'jax.debug': ['DebugEffect', 'log'],
   'jax.distributed': ['is_initialized'],
-  'jax.dlpack': ['jax'],
   'jax.dtypes': ['extended', 'finfo', 'iinfo'],
   'jax.errors': ['JAXIndexError', 'JAXTypeError'],
   'jax.ffi': ['build_ffi_lowering_function', 'include_dir', 'register_ffi_target_as_batch_partitionable', 'register_ffi_type_id'],