From 11e303a98e0bfa735e66ed116a7e7821de97315c Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Wed, 8 Jul 2026 12:24:11 -0700 Subject: [PATCH] Add pytype ignore statements related to deprecated JAX APIs PiperOrigin-RevId: 944636622 --- learned_optimization/jax_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/learned_optimization/jax_utils.py b/learned_optimization/jax_utils.py index e6a95b3..d6077b2 100644 --- a/learned_optimization/jax_utils.py +++ b/learned_optimization/jax_utils.py @@ -27,7 +27,7 @@ from jax.extend.core import unsafe_am_i_under_a_jit_DO_NOT_USE # pylint: disable=g-import-not-at-top except ImportError: # JAX v0.9.2 or older - from jax.core import unsafe_am_i_under_a_jit_DO_NOT_USE # pylint: disable=g-import-not-at-top + from jax.core import unsafe_am_i_under_a_jit_DO_NOT_USE # pylint: disable=g-import-not-at-top # pytype: disable=import-error def maybe_static_cond(pred, true_fn, false_fn, val):