From adbffc6a103e8a065ec8464f8e2bf73fbde2d26d Mon Sep 17 00:00:00 2001 From: George Necula Date: Fri, 29 Mar 2024 05:23:55 -0700 Subject: [PATCH] Replace internal usage of stop_outfeed_receiver with _deprecated_stop_outfeed_receiver. The jax.experimental.host_callback module is deprecated and will be removed. See https://github.com/google/jax/issues/20385. PiperOrigin-RevId: 620220346 --- tests/infeed_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/infeed_test.py b/tests/infeed_test.py index 37592d52fa49..572920fa4d3b 100644 --- a/tests/infeed_test.py +++ b/tests/infeed_test.py @@ -78,7 +78,7 @@ def f(x): @jax.numpy_rank_promotion("allow") # Test explicitly exercises implicit rank promotion. def testInfeedThenOutfeed(self): - hcb.stop_outfeed_receiver() + hcb._deprecated_stop_outfeed_receiver() @jax.jit def f(x): @@ -100,7 +100,7 @@ def f(x): self.assertAllClose(out, y + np.float32(1)) def testInfeedThenOutfeedInALoop(self): - hcb.stop_outfeed_receiver() + hcb._deprecated_stop_outfeed_receiver() def doubler(_, token): y, token = lax.infeed(