code: def split(
      tensor,
      split_dimension,
      num_devices,
      assign_tuple_sharding=False,
      use_sharding_op=False,
      input_shape=None,
  ):
      """Returns a tensor that is split along the given dimension.
  
      Args:
        tensor: A tf.Tensor to split.
        split_dimension: The dimension to split.
        num_devices: The number of devices to partition the dimension.
        assign_tuple_sharding: If the sharding type should be a tuple.
        use_sharding_op: If true, adds a sharding op to set the sharding.
        input_shape: The full shape of the input tensor.
      """
      return Sharding.split(
          tensor, split_dimension, num_devices, input_shape
      ).apply_to_tensor(
          tensor,
          assign_tuple_sharding=assign_tuple_sharding,
          use_sharding_op=use_sharding_op,
      )
  
imports: import numpy
  from tensorflow.compiler.tf2xla.python import xla
  from local_xla.xla import xla_data_pb2
  from tensorflow.core.framework import attr_value_pb2
label_desc: Use image node_id 5 for calling a global function with example usage: split(tensor, split_dimension, num_devices, assign_tuple_sharding, use_sharding_op, input_shape) and returns: Sharding