Function StrideOrInflate

tensorflow/core/kernels/einsum_op.cc:342–412

// the input (e.g. iii->i) or inflates it (e.g. i->iii), respectively.
template <typename Device, typename T>
Status StrideOrInflate(OpKernelContext* ctx, const Tensor& input,
                       const Labels& labels, const LabelCounts& label_counts,
                       const bool should_inflate, Tensor* output) {
  // Return early if there are no repeated indices.
  if (absl::c_all_of(label_counts, [](int c) { return c <= 1; })) {
    return CopyFrom(input, input.shape(), output);
  }
  // We reshape so that each repeated label is compressed to one dimension.
  // E.g. For iiij -> ij, The shape [3, 3, 3, 5] would be compressed to [27, 5].
  // Striding appropriately (in this case with strides 14 (=1+3+9) and 1)
  // recovers the generalized diagonal of shape [3, 5].
  ShapeVec reshape;
  ShapeVec strides;
  // Strided and inflated shapes correspond to input and output shapes,
  // respectively, should_inflate is true (vice-versa if should_inflate is
  // false). E.g. they are [3, 5] and [3, 3, 3, 5] in the above example.
  ShapeVec strided_shape;
  ShapeVec inflated_shape;
  for (int label : labels) {
    const int count = label_counts[label];
    const int current_axis =
        should_inflate ? strided_shape.size() : inflated_shape.size();
    const int64 dim = input.dim_size(current_axis);
    strided_shape.push_back(dim);
    inflated_shape.insert(inflated_shape.end(), count, dim);
    const int64 reshape_dim = MathUtil::IPow(dim, count);
    reshape.push_back(reshape_dim);
    // While taking the d-diagonal in a rank k Tensor, we take d equally-spaced
    // elements including the first and last element. Then,
    // (k - 1) * stride = d^k - 1, or, stride = (d^k - 1)/(d - 1).
    const int64 stride =
        (dim > 1 && count > 1) ? (reshape_dim - 1) / (dim - 1) : 1;
    strides.push_back(stride);
  }

  TensorShape output_shape =
      TensorShape(should_inflate ? inflated_shape : strided_shape);
  TF_RETURN_IF_ERROR(
      ctx->allocate_temp(DataTypeToEnum<T>::value, output_shape, output));
  const Device& device = ctx->eigen_device<Device>();
  switch (reshape.size()) {
#define NDIMS_CASE(N)                                                 \
  case N: {                                                           \
    if (should_inflate) {                                             \
      auto output_map = output->shaped<T, N>(reshape);                \
      auto input_map = input.shaped<T, N>(strided_shape);             \
      functor::InflateFunctor<Device, T, N>()(                        \
          device, input_map, TensorShape(strides).AsEigenDSizes<N>(), \
          output_map);                                                \
    } else {                                                          \
      auto input_map = input.shaped<T, N>(reshape);                   \
      auto output_map = output->shaped<T, N>(strided_shape);          \
      functor::StrideFunctor<Device, T, N>()(                         \
          device, input_map, TensorShape(strides).AsEigenDSizes<N>(), \
          output_map);                                                \
    }                                                                 \
  } break;
    NDIMS_CASE(1);