// CHECK-NEXT: %[[SHAPE:.*]] = "tf.RiscShape"(%arg0) {T = i32} : (tensor<2x3xf32>) -> tensor<*xi32>
// CHECK-NEXT: %[[ALPHA1:.*]] = "tf.RiscBroadcast"(%[[ALPHA]], %[[SHAPE]]) : (tensor<f32>, tensor<*xi32>) -> tensor<*xf32>
// CHECK-NEXT: %[[MAX:.*]] = "tf.RiscMaximum"(%arg0, %[[ALPHA1]]) : (tensor<2x3xf32>, tensor<*xf32>) -> tensor<*xf32>

  // CHECK: %[[CUSTOM_AGGREGATOR_2:.*]], {{.*}}, {{.*}}, {{.*}} = "tf.CustomAggregator"(%[[XLA_CALL_MODULE:.*]])

// Given a tensor of `shape` and `type`, as what shape should it be stored on
// the TPU device? This function tranposes or flattens the excessively-padded
// tensors to rank 1, but leaves other tensor shapes alone.
absl::StatusOr<xla::Shape> TpuShapeRepresentation(
    const TensorShape& shape, DataType type, bool use_fast_memory,
    XlaLayoutPreference layout_preference) {
  xla::Shape xla_shape;
  TF_RETURN_IF_ERROR(

  // Check if options have the supplied shape determination functions set.
  TF_ASSERT_OK_AND_ASSIGN(
      auto shape, options.shape_determination_fns.shape_representation_fn(
                      TensorShape(), DT_FLOAT, false,
                      tensorflow::XlaLayoutPreference::kTpuPreferLinearLayout));
  EXPECT_EQ(shape, xla::Shape());
  EXPECT_EQ(options.shape_determination_fns.layout_preference_fn(

          type: "list(shape)"
          has_minimum: true
        }
      }
      op {
        name: "Placeholder"
        output_arg {
          name: "output"
          type_attr: "dtype"
        }
        attr {
          name: "dtype"
          type: "type"
        }
        attr {
          name: "shape"
          type: "shape"
          default_value {

  auto c = ops::Add(root.WithOpName("C"), a, b);
  auto d = ops::Neg(root.WithOpName("D"), c);

  a.node()->AddAttr("_index", 0);
  b.node()->AddAttr("_index", 1);

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  TF_CHECK_OK(root.ToGraph(graph.get()));

  std::map<int, InferredShape> arg_shapes;
  arg_shapes[0].shape = TensorShape({2, 3});
  arg_shapes[1].shape = TensorShape({2, 3});

 protected:
  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.

      TensorShape shape;
      TF_RETURN_IF_ERROR(XLAShapeToTensorShape(subshape, &shape));
      output_tensor_shapes.push_back(shape);
    }
  } else {
    for (int i = 0; i < ctx->num_outputs(); ++i) {
      output_tensor_shapes.push_back(compilation_result->outputs[i].shape);
    }
  }

  // Copy XLA results to the OpOutputList.
  int output_num = 0;
  for (int i = 0, end = ctx->num_outputs(); i < end; ++i) {

// CHECK-NEXT:       shape: [ 1, 528 ],
// CHECK-NEXT:       type: INT8,
// CHECK-NEXT:       buffer: 1,
// CHECK-NEXT:       name: "arg0",
// CHECK-NEXT:       quantization: {
// CHECK-NEXT:         scale: [ 0.037248 ],
// CHECK-NEXT:         zero_point: [ -19 ]
// CHECK-NEXT:       },
// CHECK-NEXT:       has_rank: true
// CHECK-NEXT:     }, {
// CHECK-NEXT:       shape: [ 2048, 528 ],
// CHECK-NEXT:       type: INT8,

  Output shape = ops::Add(s.WithOpName("shape"), shape_a, shape_b);

  Output reshape_input = ops::Placeholder(s.WithOpName("reshape_input"),
                                          DT_FLOAT, ops::Placeholder::Attrs{});
  Output reshape = ops::Reshape(s.WithOpName("reshape"), reshape_input, shape);

  AddToCluster({shape.node(), reshape.node()}, "cluster_0");