StridedArrayViewBase(ArrayRef<int64_t> shape, ArrayRef<int64_t> index,
                       int64_t axis) {
    assert(shape.size() == index.size());
    assert(axis < shape.size());
    assert(axis >= 0);
    assert(index[axis] == 0);
    offset_ = IndexToOffset(shape, index);
    stride_ = StrideForAxis(shape, axis);
    size_ = shape[axis];
  }

  // Returns the size of the 1-d slice across the tensor.

    auto axis = EinsumGetAxisFromLabel(subscripts, s);
    if (!axis.has_value()) {
      // Should never happen.
      scope.UpdateStatus(errors::Internal(
          absl::StrCat("Missing axis", absl::string_view(&s, 1))));
    } else {
      reduced_axes.push_back(*axis);
    }
  }
  // Get the corresponding dimensions for each reduced axis.
  std::vector<Output> reduced_dims_inputs;

// CHECK-NEXT: %[[CC0:.*]] = "tf.RiscConcat"(%[[ED0]], %[[ED1]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
// CHECK-NEXT: %[[ED2:.*]] = "tf.ExpandDims"(%arg2, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: %[[CC1:.*]] = "tf.RiscConcat"(%[[CC0]], %[[ED2]]) {axis = 0 : i32} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>

  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
  auto axis = Const(scope_, 0);
  auto y = Concat(scope_, xs, axis);
  TensorShape result_shape({9, 2, 5});
  RunTest(xs, {shape, shape, shape}, {y}, {result_shape});
}

TEST_F(ArrayGradTest, BroadcastToGrad) {
  TensorShape x_shape({2, 5});

                std::vector<Output>* grad_outputs) {
  int N;
  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "N", &N));
  int axis;
  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "axis", &axis));

  grad_outputs->reserve(N);
  auto grad_op = Unstack(scope, grad_inputs[0], N, Unstack::Axis(axis));
  for (const Output& o : grad_op.output) {
    grad_outputs->emplace_back(o);
  }
  return scope.status();
}

  kBinaryOp,        // Binary ops have lhs/rhs attr.
  kQuantizationOp,  // Quantization ops have input/output attr.
};

// For each op type, the following axis carries axis information:
// kDynamicRangeOp: rhs_quantization_axis will carry axis information.
// kUnaryOp: quantization_axis will carry axis information.
// kBinaryOp: Among {lhs, rhs, output}_quantization_axis, only check rhs.

    <?x?x3x2>, axis=2 => N=6
    ```
  }];

  let arguments = (ins
    quant_RealValueType:$arg,
    ElementsAttr:$layerStats,
    OptionalAttr<ElementsAttr>:$axisStats,
    OptionalAttr<I64Attr>:$axis);
  let results = (outs quant_RealValueType);
  let hasVerifier = 1;
}

def Quantization_CoupledRefOp : Quantization_Op<"coupled_ref", [SameOperandsAndResultType]> {
  let summary = [{

                                                 ElementsAttr axis_stats,
                                                 IntegerAttr axis) {
  auto stats_op = b.create<quantfork::StatisticsOp>(
      b.getUnknownLoc(), res, layer_stats, axis_stats, axis);
  res.replaceAllUsesWith(stats_op);
  stats_op.getOperation()->replaceUsesOfWith(stats_op, res);
}

              (UpdateShapeWithAxis<-1> $qtype, $old_value))),
          [(CanUpdateShapeWithAxis<-1> $qtype, $old_value)]>;

// The axis is set to 0 because the transpose is from the legalization of
// tf.conv2d and the new channel axis is the first dimension.
def ReorderTransposeDequantQuantUsedByConv :
      Pat<(TF_TransposeOp:$old_value
              (TFL_DequantizeOp (TFL_QuantizeOp $input, $qtype)), $perm),

  : Constraint<CPred<"$0.getType().isa<RankedTensorType>()">>;

// This pattern converts TensorFlow axis format to HLO axis format which
// doesn't wrap around like TensorFlow and is always positive. For this
// conversion, use the first input to get inputs rank. Other inputs need not be
// ranked.
// Defining op for `axis` is TensorFlow constant op in the pattern as during