code: def testManyEagerOps(self, size, same_value):
      if self.mesh.device_type() != "TPU":
          # For the CPU/GPU mesh, we have a shortcut that doesn't go through the
          # MLIR, but run the eager op locally and broadcast to all the devices.
          expected_cache_diff = 0
          expected_kernel_cache = 0
          expected_device_cache = 0
          expected_eager_pure_hit = 10
      else:
          # TODO(b/287529295): Remove this branch after the TPU issue is fixed.
          expected_device_cache = 0
          expected_eager_pure_hit = 0
          if same_value:
              expected_cache_diff = 1
              expected_kernel_cache = 2
          else:
              if size >= 20:
                  expected_cache_diff = 1
                  expected_kernel_cache = 2
              else:
                  expected_cache_diff = 2
                  expected_kernel_cache = 4
  
      r = range(100)
      c_layout = Layout.replicated(self.mesh, rank=2)
      values = [
          np.reshape(r[i : i + size], (4, size // 4)) for i in range(10)
      ]
      values = [
          constant_op.constant(v, dtype=dtypes.float32) for v in values
      ]
      c0 = [api.copy_to_mesh(v, c_layout) for v in values]
  
      c0 = [c0[0 if same_value else i] for i in range(10)]
      e0 = [values[0 if same_value else i] for i in range(10)]
  
      stats1 = api._dtensor_device()._get_stats()
  
      for i in range(10):
          self.assertAllEqual(array_ops.identity(c0[i]).numpy(), e0[i])
  
      gc.collect()
  
      stats2 = api._dtensor_device()._get_stats()
      diff = diff_dicts(stats2, stats1)
  
      if same_value:
          self.assertEqual(
              diff["function_manager.size"], expected_cache_diff
          )
          self.assertEqual(
              diff["eager_pure_optimization.hit"],
              expected_eager_pure_hit,
          )
          # TFRT doesn't use eager cache.
          if not test_util.is_tfrt_enabled():
              self.assertEqual(
                  diff["kernel_cache.size"], expected_kernel_cache
              )
              self.assertEqual(
                  diff["device_cache.size"], expected_device_cache
              )
      else:
          # FIXME(feyu): Update these when the leaks are fixed.
          if size >= 20:
              self.assertEqual(
                  diff["function_manager.size"], expected_cache_diff
              )
              self.assertEqual(
                  diff["eager_pure_optimization.hit"],
                  expected_eager_pure_hit,
              )
              # TFRT doesn't use eager cache.
              if not test_util.is_tfrt_enabled():
                  self.assertEqual(
                      diff["kernel_cache.size"], expected_kernel_cache
                  )
                  self.assertEqual(
                      diff["device_cache.size"], expected_device_cache
                  )
          else:
              self.assertEqual(
                  diff["function_manager.size"], expected_cache_diff
              )
              self.assertEqual(
                  diff["eager_pure_optimization.hit"],
                  expected_eager_pure_hit,
              )
              # TFRT doesn't use eager cache.
              if not test_util.is_tfrt_enabled():
                  self.assertEqual(
                      diff["kernel_cache.size"], expected_kernel_cache
                  )
                  self.assertEqual(
                      diff["device_cache.size"], expected_device_cache
                  )
  
imports: import gc
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_stateless_random_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 8 for calling the DTensorDeviceCacheTest obj's underlying member method code with example usage: obj.testManyEagerOps(size, same_value) without return types