code: def test_MaxShardSizePolicy_1D(self):
      root = module.Module()
      with ops.device("cpu:0"):
          v0 = resource_variable_ops.ResourceVariable(
              [0.0, 1.0, 2.0, 3.0], name="v0", dtype=dtypes.float32
          )
          v1 = resource_variable_ops.ResourceVariable(
              [[4], [5], [6], [7]], name="v1", dtype=dtypes.int32
          )
      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=[4]),
                  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=[4, 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: 4 bytes
      # Each element of v0/v1 is a 32 bit/4 byte value, so each variable should be
      # split into 4 shards.
      callback = sharding_policies.MaxShardSizePolicy(max_shard_size=4)
      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",
              },
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.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], var_shape=[1]).spec
      self.assertEqual(
          self.evaluate(shards[0][v0_name][slice_spec]), 0.0
      )
  
      slice_spec = V0SaveSliceInfo(var_offset=[1], var_shape=[1]).spec
      self.assertEqual(
          self.evaluate(shards[1][v0_name][slice_spec]), 1.0
      )
  
      slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[1]).spec
      self.assertEqual(
          self.evaluate(shards[2][v0_name][slice_spec]), 2.0
      )
  
      slice_spec = V0SaveSliceInfo(var_offset=[3], var_shape=[1]).spec
      self.assertEqual(
          self.evaluate(shards[3][v0_name][slice_spec]), 3.0
      )
  
      # V1
      slice_spec = V1SaveSliceInfo(
          var_offset=[0, 0], var_shape=[1, 1]
      ).spec
      self.assertEqual(
          self.evaluate(shards[4][v1_name][slice_spec]), [4]
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[1, 0], var_shape=[1, 1]
      ).spec
      self.assertEqual(
          self.evaluate(shards[5][v1_name][slice_spec]), [5]
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[2, 0], var_shape=[1, 1]
      ).spec
      self.assertEqual(
          self.evaluate(shards[6][v1_name][slice_spec]), [6]
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[3, 0], var_shape=[1, 1]
      ).spec
      self.assertEqual(
          self.evaluate(shards[7][v1_name][slice_spec]), [7]
      )
  
      # max_shard_size: 8 bytes
      # v0/v1 haven't changed, so they should now be split into 2 shards each.
      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",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "_CHECKPOINTABLE_OBJECT_GRAPH",
              },
          ],
      )
  
      # V0
      slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[2]).spec
      self.assertAllEqual(
          self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0]
      )
  
      slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[2]).spec
      self.assertAllEqual(
          self.evaluate(shards[1][v0_name][slice_spec]), [2.0, 3.0]
      )
  
      # V1
      slice_spec = V1SaveSliceInfo(
          var_offset=[0, 0], var_shape=[2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[2][v1_name][slice_spec]), [[4], [5]]
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[2, 0], var_shape=[2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[3][v1_name][slice_spec]), [[6], [7]]
      )
  
      # max_shard_size: 10 bytes
      # 10 bytes is an uneven boundary for 4 byte elements. v0/v1 should be split
      # into 2 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",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "_CHECKPOINTABLE_OBJECT_GRAPH",
              },
          ],
      )
  
      # V0
      slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[2]).spec
      self.assertAllEqual(
          self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0]
      )
  
      slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[2]).spec
      self.assertAllEqual(
          self.evaluate(shards[1][v0_name][slice_spec]), [2.0, 3.0]
      )
  
      # V1
      slice_spec = V1SaveSliceInfo(
          var_offset=[0, 0], var_shape=[2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[2][v1_name][slice_spec]), [[4], [5]]
      )
  
      slice_spec = V1SaveSliceInfo(
          var_offset=[2, 0], var_shape=[2, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[3][v1_name][slice_spec]), [[6], [7]]
      )
  
      # max_shard_size: 16 bytes
      # 16 bytes the exact size of each variable, so they should get 1 shard each.
      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",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "_CHECKPOINTABLE_OBJECT_GRAPH",
              },
          ],
      )
  
      # V0
      slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[4]).spec
      self.assertAllEqual(
          self.evaluate(shards[0][v0_name][slice_spec]),
          [0.0, 1.0, 2.0, 3.0],
      )
  
      # V1
      slice_spec = V1SaveSliceInfo(
          var_offset=[0, 0], var_shape=[4, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[1][v1_name][slice_spec]),
          [[4], [5], [6], [7]],
      )
  
      # max_shard_size: 18 bytes
      # 18 bytes slightly larger than the size of each variable, but not large
      # enough to fit another 4 byte element, so they should get 1 shard each.
      callback = sharding_policies.MaxShardSizePolicy(max_shard_size=18)
      shards = []
      for tensors in shardable_tensors:
          shards.extend(callback(tensors))
  
      self.assertEqual(
          [set(shard.keys()) for shard in shards],
          [
              {
                  "v0/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "v1/.ATTRIBUTES/VARIABLE_VALUE",
              },
              {
                  "_CHECKPOINTABLE_OBJECT_GRAPH",
              },
          ],
      )
  
      # V0
      slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[4]).spec
      self.assertAllEqual(
          self.evaluate(shards[0][v0_name][slice_spec]),
          [0.0, 1.0, 2.0, 3.0],
      )
  
      # V1
      slice_spec = V1SaveSliceInfo(
          var_offset=[0, 0], var_shape=[4, 1]
      ).spec
      self.assertAllEqual(
          self.evaluate(shards[1][v1_name][slice_spec]),
          [[4], [5], [6], [7]],
      )
  
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 3 for calling the ShardingPoliciesTest obj's underlying member method code with example usage: obj.test_MaxShardSizePolicy_1D() without return types