From 607f3c66b537ca7b3e3800f1e9e40e1f97a15c16 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Fri, 28 Oct 2022 15:19:52 -0700
Subject: [PATCH] Silenced warning (#40)

---
 jaxtyping/pytree_type.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/jaxtyping/pytree_type.py b/jaxtyping/pytree_type.py
index b97a208..4736fff 100644
--- a/jaxtyping/pytree_type.py
+++ b/jaxtyping/pytree_type.py
@@ -22,7 +22,7 @@ import typing
 from typing import Generic, TYPE_CHECKING, TypeVar
 from typing_extensions import Protocol
 
-import jax
+import jax.tree_util as jtu
 import typeguard
 
 
@@ -83,7 +83,7 @@ class _MetaSubscriptPyTree(type):
             else:
                 return True
 
-        leaves = jax.tree_leaves(obj, is_leaf=is_leaftype)
+        leaves = jtu.tree_leaves(obj, is_leaf=is_leaftype)
         return all(map(is_leaftype, leaves))