code: def test_CheckpointOption_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")
      self.evaluate(v0.initializer)
      with ops.device("/job:worker/task:1/cpu:0"):
          v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
      self.evaluate(v1.initializer)
      with ops.device("/job:worker/task:2/cpu:0"):
          v2 = resource_variable_ops.ResourceVariable(2.0, name="v2")
      self.evaluate(v2.initializer)
      root.v0 = v0
      root.v1 = v1
      root.v2 = v2
  
      tmp_dir = self.create_tempdir("ckpt")
      ckpt = checkpoint.Checkpoint(root)
      save_path = ckpt.save(
          tmp_dir,
          options=checkpoint_options.CheckpointOptions(
              experimental_sharding_callback=(
                  sharding_policies.ShardByTaskPolicy()
              )
          ),
      )
      self.assertLen(gfile.Glob(save_path + ".data*"), 4)
      ckpt.restore(save_path)
  
imports: import random
  import string
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_options
  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.platform import gfile
  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_CheckpointOption_ShardByTaskPolicy() without return types