code: def test_default_mesh(self):
      @polymorphic_function.function
      def func(a):
          return a + 3.0
  
      with api.default_mesh(self.mesh):
          a = array_ops.zeros(shape=())
          result = func(a)
  
      self.assertEqual(result, 3.0)
      self.assertEqual(api.fetch_layout(result).mesh, self.mesh)
      self.assertTrue(api.fetch_layout(result).is_fully_replicated())
      self.assertEqual(result.device, api.device_name())
  
      # Also make sure it works as wrapper
      @api.default_mesh(self.mesh)
      def func2():
          b = array_ops.ones(shape=())
          return func(b)
  
      result = func2()
      self.assertEqual(result, 4.0)
      self.assertEqual(api.fetch_layout(result).mesh, self.mesh)
      self.assertTrue(api.fetch_layout(result).is_fully_replicated())
      self.assertEqual(result.device, api.device_name())
  
      with self.assertRaisesRegex(
          ValueError, "Expect `mesh` to be `Mesh`"
      ):
          with api.default_mesh(None):
              pass
  
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_random
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 8 for calling the APITest obj's underlying member method code with example usage: obj.test_default_mesh() and returns: unknown, func