From 7f85a12a85afdba03ba146b1f62b71100d2f936c Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Sun, 15 Oct 2023 13:16:00 -0700
Subject: [PATCH] Added DTypeLike to match jax.typing

---
 jaxtyping/__init__.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index e66d044..d6ecfca 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -65,6 +65,8 @@ elif has_jax:
 if typing.TYPE_CHECKING:
     # Introduce an indirection so that we can `import X as X` to make it clear that
     # these are public.
+    from jax.typing import DTypeLike as DTypeLike
+
     from ._indirection import (
         BFloat16 as BFloat16,
         Bool as Bool,
@@ -121,8 +123,13 @@ else:
     )
 
     if has_jax:
+        import jax.typing
+
         from ._array_types import Key as Key
 
+        if hasattr(jax.typing, "DTypeLike"):
+            from jax.typing import DTypeLike as DTypeLike
+
 
 # Now import PyTreeDef and PyTree
 if typing.TYPE_CHECKING: