code: def test_manual_sharding_ops(self):
      strategy, num_replicas = get_tpu_strategy(enable_spmd=True)
  
      @def_function.function
      def fn(x):
          x_split = strategy.experimental_split_to_logical_devices(
              x, [1, 2]
          )
          split_sharding = xla_sharding.get_op_sharding(x_split.op)
          x_manual = xla_sharding.auto_to_manual_spmd_partition(
              x_split, split_sharding
          )
          y_manual = x_manual + 1
          y_split = xla_sharding.manual_to_auto_spmd_partition(
              y_manual, split_sharding, (2, 2)
          )
          return y_split
  
      arg = constant_op.constant(0, shape=(2, 2), dtype=dtypes.int64)
      result = strategy.run(fn, args=(arg,))
      self.assertAllEqual(
          (arg + 1) * num_replicas,
          self.evaluate(strategy.reduce("SUM", result, axis=None)),
      )
  
imports: import os
  from absl.testing import parameterized
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.compiler.xla.experimental import xla_sharding
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import packed_distributed_variable
  from tensorflow.python.distribute import strategy_test_lib
  from tensorflow.python.distribute import tpu_replicated_variable
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute import tpu_values
  from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import remote
  from tensorflow.python.eager import test
  from tensorflow.python.framework import config
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import summary_test_util
  from tensorflow.python.module import module
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops import summary_ops_v2
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import flags
  from tensorflow.python.tpu import device_assignment
  from tensorflow.python.tpu import tpu_replication
label_desc: Use image node_id 13 for calling the TPUStrategyModelParallelismTest obj's underlying member method code with example usage: obj.test_manual_sharding_ops() and returns: y_split