From c0cb70720d85c64a501b08d902a7a4edbdbecddc Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Thu, 6 Aug 2026 08:22:38 -0700 Subject: [PATCH] add pyrefly suppressions for jax.jit annotation change PiperOrigin-RevId: 960314905 --- examples/sst2/train_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/sst2/train_test.py b/examples/sst2/train_test.py index 6a8b0585d..2e9bea969 100644 --- a/examples/sst2/train_test.py +++ b/examples/sst2/train_test.py @@ -46,7 +46,7 @@ def test_train_step_updates_parameters(self): batch = {'token_ids': token_ids, 'length': lengths, 'label': labels} rngs = {'dropout': rng} train_step_fn = jax.jit(train.train_step) - new_state, metrics = train_step_fn(state, batch, rngs) + new_state, metrics = train_step_fn(state, batch, rngs) # pyrefly: ignore[bad-argument-type] self.assertIsInstance(new_state, train.TrainState) self.assertIsInstance(metrics, train.Metrics) old_param_values = jax.tree_util.tree_leaves(state.params)