}
      },
      sharding_or_op);
}

// Returns the sharding from a op-sharding variant if it is available and valid.
std::optional<xla::OpSharding> GetShardingFromVariant(
    const OpShardingVariant& sharding_or_op) {
  xla::OpSharding sharding;
  const auto sharding_string = GetShardingStringFromVariant(sharding_or_op);
  if (!sharding_string) return std::nullopt;