def split_to_logical_devices(tensor: TfVal, partition_dimensions: pxla.PartitionsOrReplicated):
    'Like TPUMPStrategy.experimental_split_to_logical_devices.\n\n  For jax2tf purposes we want to avoid needing to thread the `strategy` object\n  through the generated computation. It seems that the original function needs\n  the strategy object only for error checking, which we assume is done upstream\n  by JAX.\n\n  Args:\n    tensor: Input tensor to annotate.\n    partition_dimensions: A list of integers, with one integer per tensor\n      dimension, specifying in how many parts the dimension should be split. The\n      product of integers must equal the number of devices per replica.\n    use_sharding_op: whether to use a sharding op, or not.\n\n  Returns:\n    an annotated tensor.\n  '
    if (partition_dimensions is None):
        return xla_sharding.replicate(tensor, use_sharding_op=True)
    num_partition_splits = np.prod(partition_dimensions)
    tile_assignment = np.arange(num_partition_splits).reshape(partition_dimensions)
    return xla_sharding.tile(tensor, tile_assignment, use_sharding_op=True)
Like TPUMPStrategy.experimental_split_to_logical_devices.

For jax2tf purposes we want to avoid needing to thread the `strategy` object
through the generated computation. It seems that the original function needs
the strategy object only for error checking, which we assume is done upstream
by JAX.

Args:
  tensor: Input tensor to annotate.
  partition_dimensions: A list of integers, with one integer per tensor
    dimension, specifying in how many parts the dimension should be split. The
    product of integers must equal the number of devices per replica.
  use_sharding_op: whether to use a sharding op, or not.

Returns:
  an annotated tensor.
def split_to_logical_devices(tensor: TfVal, partition_dimensions: pxla.PartitionsOrReplicated):
    'Like TPUMPStrategy.experimental_split_to_logical_devices.\n\n  For jax2tf purposes we want to avoid needing to thread the `strategy` object\n  through the generated computation. It seems that the original function needs\n  the strategy object only for error checking, which we assume is done upstream\n  by JAX.\n\n  Args:\n    tensor: Input tensor to annotate.\n    partition_dimensions: A list of integers, with one integer per tensor\n      dimension, specifying in how many parts the dimension should be split. The\n      product of integers must equal the number of devices per replica.\n    use_sharding_op: whether to use a sharding op, or not.\n\n  Returns:\n    an annotated tensor.\n  '
    if (partition_dimensions is None):
        return xla_sharding.replicate(tensor, use_sharding_op=True)
    num_partition_splits = np.prod(partition_dimensions)
    tile_assignment = np.arange(num_partition_splits).reshape(partition_dimensions)
    return xla_sharding.tile(tensor, tile_assignment, use_sharding_op=True)