// the full shape is (num_cores_per_replica * a, b, c), then the shard shape is
// (a, b, c). `context_op` is used for error reporting, in case of errors.
LogicalResult GetShardShapedType(Operation* context_op,
                                 int num_cores_per_replica, Type full_type,
                                 Type& shard_type) {