code: def test_computation_follows_data(self):
      if jax.device_count() < 5:
          self.skipTest("test requires 5 devices")
      devices = self.get_devices()
  
      # By default, computation is placed (uncommitted) on device 0
      x = jnp.ones(2)
      self.assert_uncommitted_to_device(x, devices[0])
  
      # Computing only with uncommitted data, will perform the computation
      # on the same device 0, and the result is uncommitted
      y = jnp.sin(x + x * 1) + 2
      self.assert_uncommitted_to_device(y, devices[0])
  
      # We can use device_put to both place the data on the requested device,
      # and to commit it there.
      z = jax.device_put(1, devices[1])
      self.assert_committed_to_device(z, devices[1])
  
      # A computation with some committed inputs will happen on the committed
      # device. Uncommitted data may be moved to the committed device.
      u = z + x  # z is committed, and x is uncommitted
      self.assert_committed_to_device(u, devices[1])
  
      # A computation with inputs committed to multiple devices will result
      # in an error
  
      err_msg = "Received incompatible devices for jitted computation"
      with self.assertRaisesRegex(ValueError, err_msg):
          jax.device_put(x, devices[2]) + jax.device_put(x, devices[3])
  
      # A jitted-computation without a device specification behave like any
      # other primitive
      jit_add = jax.jit(lambda a, b: a + b)
      self.assert_uncommitted_to_device(jit_add(1, 2), devices[0])
      self.assert_committed_to_device(
          jit_add(1, jax.device_put(2, devices[2])), devices[2]
      )
      with self.assertRaisesRegex(ValueError, err_msg):
          jit_add(
              jax.device_put(x, devices[2]),
              jax.device_put(x, devices[3]),
          )
  
      # Even jit of trivial computations leaves the result uncommitted
      x_uncommitted = jnp.array([1, 2, 3])
      y = jax.jit(lambda x: x)(x_uncommitted)
      self.assert_uncommitted_to_device(y, devices[0])
  
      z1, z2 = jax.jit(lambda x: (x, x))(x_uncommitted)
      self.assert_uncommitted_to_device(z1, devices[0])
      self.assert_uncommitted_to_device(z2, devices[0])
      # trivial computation does not exist in JAX anymore.
      self.assertNotEqual(
          z1.unsafe_buffer_pointer(), z2.unsafe_buffer_pointer()
      )
  
      x2_uncommitted = jnp.array([2, 3])
      z1, z2, z3 = jax.jit(lambda x, y: (y, 1, x))(
          x_uncommitted, x2_uncommitted
      )
      self.assert_uncommitted_to_device(z1, devices[0])
      self.assert_uncommitted_to_device(z2, devices[0])
      self.assert_uncommitted_to_device(z3, devices[0])
  
      # A jitted computation with an device specification behaves as if the
      # arguments are first device_put to the specified device. The result
      # will be committed on the specified.
      # The `device` parameter is experimental, and subject to change.
      jit_add_on4 = jax.jit(lambda a, b: a + b, device=devices[4])
      self.assert_committed_to_device(jit_add_on4(1, 2), devices[4])
      self.assert_committed_to_device(
          jit_add_on4(1, jax.device_put(2, devices[2])), devices[4]
      )
      self.assert_committed_to_device(
          jit_add_on4(
              jax.device_put(x_uncommitted, devices[2]),
              jax.device_put(x_uncommitted, devices[3]),
          ),
          devices[4],
      )
  
imports: import os
  from unittest import SkipTest
  from absl.testing import absltest
  import jax
  import jax.numpy
  from jax import lax
  from jax._src import test_util
  from jax._src import xla_bridge
  from jax import config
label_desc: Use image node_id 4 for calling the MultiDeviceTest obj's underlying member method code with example usage: obj.test_computation_follows_data() without return types