code: def test_MaxShardSizePolicy_2D(self):
      root = module.Module()
      with ops.device("cpu:0"):
          v0 = resource_variable_ops.ResourceVariable(
              [[0, 1], [2, 3], [4, 5]], name="v0"
          )
          v1 = resource_variable_ops.ResourceVariable(
              [[[6.0], [7.0]], [[8.0], [9.0]], [[10.0], [11.0]]],
              name="v1",
          )
      self.evaluate(v0.initializer)
      self.evaluate(v1.initializer)
      root.v0 = v0
      root.v1 = v1
  
      v0_name = "v0/.ATTRIBUTES/VARIABLE_VALUE"
      v1_name = "v1/.ATTRIBUTES/VARIABLE_VALUE"
  
      class V0SaveSliceInfo(variables.Variable.SaveSliceInfo):
          def __init__(self, var_offset, var_shape):
              super().__init__(
                  full_name=v0_name,
                  full_shape=tensor_shape.TensorShape(dims=[3, 2]),
                  var_offset=var_offset,
                  var_shape=var_shape,
              )
  
      class V1SaveSliceInfo(variables.Variable.SaveSliceInfo):
          def __init__(self, var_offset, var_shape):
              super().__init__(
                  full_name=v1_name,
                  full_shape=tensor_shape.TensorShape(dims=[3, 2, 1]),
                  var_offset=var_offset,
                  var_shape=var_shape,
              )
  
      shardable_tensors = self._get_shardable_tensors_by_task(root)
  
      # Test sharding the v0 & v1 tensors with different max shard sizes.
  
      # max_shard_size: 8 bytes
      # Each element of v0/v1 is a 32 bit/4 byte value, so each variable should be
      # split into 3 shards.
      callback = sharding_policies.MaxShardSizePolicy(max_shard_size=8)
      shards = []
      for tensors in shardable_tensors:
          shards.extend(callback(tensors))
  
      self.assertEqual(
          [set(shard.keys()) for shard in shards],
          [
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "_CHECKPOINTABLE_OBJECT_GRAPH",
              },
          ],
      )
  
      # V0
      slice_spec = V0SaveSliceInfo(
          var_offset=[0, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[0][v0_name][slice_spec]), [[0, 1]]
      )
  
      slice_spec = V0SaveSliceInfo(
          var_offset=[1, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[1][v0_name][slice_spec]), [[2, 3]]
      )
  
      slice_spec = V0SaveSliceInfo(
          var_offset=[2, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[2][v0_name][slice_spec]), [[4, 5]]
      )
  
      # V1
      slice_spec = V1SaveSliceInfo(
          var_offset=[0, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[3][v1_name][slice_spec]),
          [[[6.0], [7.0]]],
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[1, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[4][v1_name][slice_spec]),
          [[[8.0], [9.0]]],
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[2, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[5][v1_name][slice_spec]),
          [[[10.0], [11.0]]],
      )
  
      # max_shard_size: 10 bytes
      # 10 bytes is an uneven boundary for 4 byte elements. v0/v1 should be split
      # into 3 shards each.
      callback = sharding_policies.MaxShardSizePolicy(max_shard_size=10)
      shards = []
      for tensors in shardable_tensors:
          shards.extend(callback(tensors))
  
      self.assertEqual(
          [set(shard.keys()) for shard in shards],
          [
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "_CHECKPOINTABLE_OBJECT_GRAPH",
              },
          ],
      )
  
      # V0
      slice_spec = V0SaveSliceInfo(
          var_offset=[0, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[0][v0_name][slice_spec]), [[0, 1]]
      )
  
      slice_spec = V0SaveSliceInfo(
          var_offset=[1, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[1][v0_name][slice_spec]), [[2, 3]]
      )
  
      slice_spec = V0SaveSliceInfo(
          var_offset=[2, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[2][v0_name][slice_spec]), [[4, 5]]
      )
  
      # V1
      slice_spec = V1SaveSliceInfo(
          var_offset=[0, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[3][v1_name][slice_spec]),
          [[[6.0], [7.0]]],
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[1, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[4][v1_name][slice_spec]),
          [[[8.0], [9.0]]],
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[2, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[5][v1_name][slice_spec]),
          [[[10.0], [11.0]]],
      )
  
      # max_shard_size: 12 bytes
      # 12 bytes is enough to fit 3 elements per variable in each shard, BUT that
      # would require concurrent multidimensional tensor partitioning, which is
      # not currently implemented for MaxShardSizePolicy. (When partitioning a
      # tensor into a shard, we choose an axis to partition along. This can
      # happen multiple times for a given tensor (in the case that the tensor
      # spans multiple shards). In that case, multiple dimensions can be
      # partitioned along (each time the tensor is partitioned, a new axis can be
      # chosen), but not within a single iteration of adding a tensor partition to
      # the shard.) So, v0/v1 should be split into 3 shards each.
      callback = sharding_policies.MaxShardSizePolicy(max_shard_size=12)
      shards = []
      for tensors in shardable_tensors:
          shards.extend(callback(tensors))
  
      self.assertEqual(
          [set(shard.keys()) for shard in shards],
          [
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "_CHECKPOINTABLE_OBJECT_GRAPH",
              },
          ],
      )
  
      # V0
      slice_spec = V0SaveSliceInfo(
          var_offset=[0, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[0][v0_name][slice_spec]), [[0, 1]]
      )
  
      slice_spec = V0SaveSliceInfo(
          var_offset=[1, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[1][v0_name][slice_spec]), [[2, 3]]
      )
  
      slice_spec = V0SaveSliceInfo(
          var_offset=[2, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[2][v0_name][slice_spec]), [[4, 5]]
      )
  
      # V1
      slice_spec = V1SaveSliceInfo(
          var_offset=[0, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[3][v1_name][slice_spec]),
          [[[6.0], [7.0]]],
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[1, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[4][v1_name][slice_spec]),
          [[[8.0], [9.0]]],
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[2, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[5][v1_name][slice_spec]),
          [[[10.0], [11.0]]],
      )
  
      # max_shard_size: 16 bytes
      # Each variable should be split into 1.5 shards. The middle shard will
      # contain elements from both variables.
      callback = sharding_policies.MaxShardSizePolicy(max_shard_size=16)
      shards = []
      for tensors in shardable_tensors:
          shards.extend(callback(tensors))
  
      self.assertEqual(
          [set(shard.keys()) for shard in shards],
          [
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "_CHECKPOINTABLE_OBJECT_GRAPH",
              },
          ],
      )
  
      # V0
      slice_spec = V0SaveSliceInfo(
          var_offset=[0, 0], var_shape=[2, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[0][v0_name][slice_spec]),
          [[0, 1], [2, 3]],
      )
  
      slice_spec = V0SaveSliceInfo(
          var_offset=[2, 0], var_shape=[1, 2]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[1][v0_name][slice_spec]), [[4, 5]]
      )
  
      # V1
      slice_spec = V1SaveSliceInfo(
          var_offset=[0, 0, 0], var_shape=[1, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[1][v1_name][slice_spec]),
          [[[6.0], [7.0]]],
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[1, 0, 0], var_shape=[2, 2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[2][v1_name][slice_spec]),
          [[[8.0], [9.0]], [[10.0], [11.0]]],
      )
  
imports: import random
  import string
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import graph_view
  from tensorflow.python.checkpoint.sharding import sharding_policies
  from tensorflow.python.checkpoint.sharding import sharding_util
  from tensorflow.python.eager import remote
  from tensorflow.python.eager import test
  from tensorflow.python.framework import device
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import test_util
  from tensorflow.python.module import module
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.training import server_lib
  from tensorflow.python.training.saving import saveable_object
  from tensorflow.python.training.saving import saveable_object_util
label_desc: Use image node_id 4 for calling the ShardingPoliciesTest obj's underlying member method code with example usage: obj.test_MaxShardSizePolicy_2D() without return types