// Returns number of elements across all dimensions.
  virtual Status NumElements(int64_t* num_elements) const = 0;
  // Returns size of specified dimension
  //
  // -1 indicates an unknown axis length; this is unreachable for most standard
  // ImmediateExecutionTensorHandles, but comes up for example when computing
  // the shape of a parallel tensor with component shapes differing across
  // devices.

}

// -----

// CHECK-LABEL: func @concat_uniform_quantize
func.func @concat_uniform_quantize(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>) -> tensor<6x3x!tf_type.qint8> {
  %axis = "tf.Const"() { value = dense<0> : tensor<i64> } : () -> tensor<i64>
  %scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
  %zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>

}

// Op: Sum()
// Summary: Computes the sum of elements across dimensions of a tensor.
//
// Description:
//   Reduces `input` along the dimensions given in `axis`. Unless
//   `keep_dims` is true, the rank of the tensor is reduced by 1 for each entry
//   in `axis`. If `keep_dims` is true, the reduced dimensions are retained with
//   length 1.
Status Sum(AbstractContext* ctx, AbstractTensorHandle* const input,

// CHECK: "tfl.unpack"(%arg0) <{axis = 1 : i32, num = 3 : i32}> : (tensor<2x3xi32>) -> (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>)
}

func.func @unpackNegAxis(%arg0: tensor<2x3xi32>) -> tensor<2xi32> {
  %0:3 = "tf.Unpack"(%arg0) {axis = -1 : i64} : (tensor<2x3xi32>) -> (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>)
  func.return %0#0 : tensor<2xi32>

// CHECK-LABEL: unpackNegAxis

    // expected-remark@+1 {{Cost: 5}}
    %11 = "tf.Pack"(%9, %0, %10, %2) {axis = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<4xi32>
    // expected-remark@+1 {{Cost: 1}}

                         %weight : tensor<*xi8>, %input : tensor<*xi32>, %axis : tensor<i32>,
                         %weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
      attributes {tf_quant.quantized_ops = ${quantized_ops}}
  {
    %accum_out = "tf.GatherV2"(%weight, %input, %axis) {

  bool reverse;
  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "reverse", &reverse));
  attrs.reverse_ = !reverse;

  auto axis = op.input(1);
  auto sum = Cumsum(scope, grad_inputs[0], axis, attrs);
  grad_outputs->push_back(sum.out);
  grad_outputs->push_back(NoGradient());
  return scope.status();
}
REGISTER_GRADIENT_OP("Cumsum", CumsumGrad);

  %2 = arith.constant dense<[]> : tensor<0xi32>
  %3 = "tfl.concatenation"(%0, %1, %2) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<2xi32>, tensor<2xi32>, tensor<0xi32>) -> tensor<?xi32>
  func.return %3 : tensor<?xi32>

  // CHECK: %0 = "tfl.concatenation"(%[[CST]], %[[CST]]) <{axis = 0 : i32, fused_activation_function = "NONE"}>
  // CHECK: return %0 : tensor<?xi32>
}

  %2 = "tfl.add"(%arg0, %arg3) {fused_activation_function = "RELU6", per_device_costs = {CPU = 5.0 : f32, GPU = 1.0 : f32}, tac.device = "GPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  %3 = "tfl.pack"(%1, %2) {axis = 0 : i32, per_device_costs = {CPU = 2.0 : f32, GPU = -1.0 : f32}, values_count = 2 : i32, tac.device = "CPU"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
  func.return %3 : tensor<2x1xf32>
})";

    return type_

  def _pack_tensor_list(self, value):
    # This is packing a list of tensors, then the axis is 0.
    axis = self._ssa_name('zero')
    self._emit_with_loc('\n{} = arith.constant 0 : i64'.format(axis))
    casted = self._ssa_name('pack')
    self.emit('\n{} = tfr.call @tf__pack({}, {})'.format(casted, value, axis))
    self._emit_with_loc(' : (!tfr.tensor_list, i64) -> !tfr.tensor')
    # load the op def of tf.Pack