// Split nodes at ith depth from the original input node represent nodes
  // that split the input data at i-th dimension.
  auto dimension_to_splits_map =
      GetDimensionIndicesAndNumSplitsFromSharding(input_sharding);
  if (!dimension_to_splits_map.ok()) {
    LOG(ERROR) << dimension_to_splits_map.status();
    return mlir::failure();
  }

bool IsReplicatedSharding(const xla::OpSharding& sharding);

// Returns a map of dimension indices and number of splits for tiled sharding.
absl::StatusOr<std::map<int, int>> GetDimensionIndicesAndNumSplitsFromSharding(
    const xla::OpSharding& sharding);
}  // namespace tensorflow