Function Contract

tensorflow/compiler/tf2xla/lib/data_format.cc:24–48

namespace {

xla::StatusOr<xla::XlaOp> Contract(xla::XlaOp input, int64 dim) {
  xla::XlaBuilder* builder = input.builder();
  TF_ASSIGN_OR_RETURN(xla::Shape input_shape, builder->GetShape(input));

  if (input_shape.dimensions().back() != 4) {
    return errors::InvalidArgument("Expected last dimension to be 4; got ",
                                   input_shape.dimensions().back());
  }

  // Transpose the input so C is directly followed by VECT_C.
  std::vector<int64> permutation;
  for (int64 i = 0; i != input_shape.rank() - 1; ++i) {
    permutation.push_back(i);
    if (i == dim) {
      permutation.push_back(input_shape.rank() - 1);
    }
  }

  // Now merge the adjacent dimensions with a reshape.
  std::vector<int64> contracted_shape(input_shape.dimensions().begin(),
                                      input_shape.dimensions().end() - 1);
  contracted_shape[dim] *= 4;

  return xla::Reshape(xla::Transpose(input, permutation), contracted_shape);
}

xla::StatusOr<xla::XlaOp> Expand(xla::XlaOp input, int64 dim) {
  xla::XlaBuilder* builder = input.builder();