code: def test_ShardByTaskPolicy(self):
      servers = [
          server_lib.Server.create_local_server() for _ in range(3)
      ]
      cluster_spec = server_lib.ClusterSpec(
          {"worker": [s.target[len("grpc://") :] for s in servers]}
      )
      remote.connect_to_cluster(cluster_spec)
      root = module.Module()
      with ops.device("/job:worker/task:0/cpu:0"):
          v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
      with ops.device("/job:worker/task:1/cpu:0"):
          v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
      with ops.device("/job:worker/task:2/cpu:0"):
          v2 = resource_variable_ops.ResourceVariable(2.0, name="v2")
      root.v0 = v0
      root.v1 = v1
      root.v2 = v2
  
      shardable_tensors = self._get_shardable_tensors_by_task(root)
  
      callback = sharding_policies.ShardByTaskPolicy()
      shards = []
      for tensors in shardable_tensors:
          shards.extend(callback(tensors))
  
      self.assertAllEqual(
          [set(shard.keys()) for shard in shards],
          [
              {"v0/.ATTRIBUTES/VARIABLE_VALUE"},
              {"v1/.ATTRIBUTES/VARIABLE_VALUE"},
              {"v2/.ATTRIBUTES/VARIABLE_VALUE"},
              {"_CHECKPOINTABLE_OBJECT_GRAPH"},
          ],
      )
  
      self.assertEqual(
          self.evaluate(shards[0]["v0/.ATTRIBUTES/VARIABLE_VALUE"][""]),
          v0.numpy(),
      )
      self.assertEqual(
          self.evaluate(shards[1]["v1/.ATTRIBUTES/VARIABLE_VALUE"][""]),
          v1.numpy(),
      )
      self.assertEqual(
          self.evaluate(shards[2]["v2/.ATTRIBUTES/VARIABLE_VALUE"][""]),
          v2.numpy(),
      )
  
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 2 for calling the ShardingPoliciesTest obj's underlying member method code with example usage: obj.test_ShardByTaskPolicy() without return types