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();
}

  // CHECK-DAG: %[[ITEMS1_3:.*]] = "tf.ExpandDims"(%[[ITEMS1]]#3, %[[AXIS]])
  // CHECK-DAG: %[[ITEMS1_2:.*]] = "tf.ExpandDims"(%[[ITEMS1]]#2, %[[AXIS]])
  // CHECK-DAG: %[[ITEMS1_1:.*]] = "tf.ExpandDims"(%[[ITEMS1]]#1, %[[AXIS]])
  // CHECK-DAG: %[[ITEMS1_0:.*]] = "tf.ExpandDims"(%[[ITEMS1]]#0, %[[AXIS]])
  // CHECK-DAG: %[[ITEMS0_0:.*]] = "tf.ExpandDims"(%[[ITEMS0]], %[[AXIS]])

      @Override
      public int size() {
        return axes.size();
      }

      @Override
      public E get(int axis) {
        checkElementIndex(axis, size());
        int axisIndex = getAxisIndexForProductIndex(index, axis);
        return axes.get(axis).get(axisIndex);
      }

      @Override
      boolean isPartialView() {
        return true;
      }

}

inline bool IsUnknownDimOrRank(int64_t dim_or_rank) {
  return dim_or_rank == -1;
}

// Returns dimension index for the given TensorFlow axis that supports negative
// indexing.
inline int64_t GetDimForAxis(int64_t axis, int64_t rank) {
  return axis >= 0 ? axis : axis + rank;
}

// Returns the tf.Equal/tf.NotEqual result type given `x` and `y` and inputs. If

  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.

module {
func.func @main(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6x3xf32> {
  %axis = "tf.Const"() { value = dense<0> : tensor<i64> } : () -> tensor<i64>
  %1 = "tf.ConcatV2"(%arg0, %arg1, %axis) : (tensor<3x3xf32>, tensor<3x3xf32>, tensor<i64>) -> tensor<6x3xf32>
  func.return %1 : tensor<6x3xf32>
}

    <?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 = [{

                $.each(plot.getXAxes(), function(_, axis) {
                    const opts = axis.options;
                    opts.min = reset ? null : ranges.xaxis.from;
                    opts.max = reset ? null : ranges.xaxis.to;
                });
                $.each(plot.getYAxes(), function(_, axis) {
                    const opts = axis.options;
                    opts.min = reset ? 0 : ranges.yaxis.from;

                                                 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),