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 #20385. PiperOrigin-RevId: 620220346
jax-ml · Mar 29, 2024 · adbffc6 · adbffc6
1 parent 15e2e26
commit adbffc6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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(