code: def testInfeedThenOutfeedInALoop(self):
      hcb.stop_outfeed_receiver()
  
      def doubler(_, token):
          y, token = lax.infeed(
              token, shape=core.ShapedArray((3, 4), jnp.float32)
          )
          return lax.outfeed(token, y * np.float32(2))
  
      @jax.jit
      def f(n):
          token = lax.create_token(n)
          token = lax.fori_loop(0, n, doubler, token)
          return n
  
      device = jax.local_devices()[0]
      n = 10
      execution = threading.Thread(target=lambda: f(n))
      execution.start()
      for _ in range(n):
          x = self.rng().randn(3, 4).astype(np.float32)
          device.transfer_to_infeed((x,))
          (y,) = device.transfer_from_outfeed(
              xla_client.shape_from_pyval(
                  (x,)
              ).with_major_to_minor_layout_if_absent()
          )
          self.assertAllClose(y, x * np.float32(2))
      execution.join()
  
file_path: jax/tests/infeed_test.py
imports: import threading
  from unittest import SkipTest
  from absl.testing import absltest
  import jax
  from jax import lax, numpy
  from jax import config
  from jax.experimental import host_callback
  from jax._src import core
  from jax._src import xla_bridge
  from jax._src.lib import xla_client
  import jax._src.test_util
  import numpy
label_desc: Use image node_id 5 for calling the InfeedTest obj's underlying member method code with example usage: obj.testInfeedThenOutfeedInALoop() and returns: lax, n