};
  TF_EXPECT_OK(ShapeAnnotationsMatch(graph, shape_info, expected));
}

TEST(ShapeInferenceTest, WhileLoopWithResource) {
  // Graph:
  // x = resource_variable_ops.var_handle_op(dtype=dtypes.float32, shape=[2, 3])
  // y = control_flow_ops.while_loop(lambda _: true, lambda x: x, [x])
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

    auto x =

      "SavedModel checkpoint does not contain object graph.");
  if (object_graph_tensor.dtype() != DT_STRING ||
      object_graph_tensor.dims() != 0 ||
      object_graph_tensor.NumElements() != 1) {
    return absl::Status(
        absl::StatusCode::kFailedPrecondition,
        "SavedModel checkpoint object graph was not the correct type.");
  }

  const tstring* object_graph_string = reinterpret_cast<const tstring*>(

      gtl::ArraySlice<AbstractTensorHandle*> output_gradients,
      absl::Span<AbstractTensorHandle*> result) const override;

  // Builds a tensor filled with ones with the same shape and dtype as `t`.
  Status BuildOnesLike(const TapeTensor& t,
                       AbstractTensorHandle** result) const override;

  // Looks up the ID of a Gradient.

      // Args
      {"x: T"},
      // Return values
      {"y: T"},
      // Attr def
      {"T: {float, double, int32, int64}"},
      // Nodes
      {
          {{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
          {{"scale"}, "Cast", {"two"}, {{"SrcT", DT_INT64}, {"DstT", "$T"}}},
          {{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}},
      });
}

class MlirGraphOptimizationPassTest : public Test {

    auto iter = placeholders.find(placeholder_index);
    Node* placeholder_node;
    if (iter == placeholders.end()) {
      NodeDefBuilder placeholder_builder(new_name, "Placeholder");
      placeholder_builder.Attr("dtype", src->output_type(src_output));
      string outside_compilation_attr;
      TF_RETURN_IF_ERROR(GetNodeAttr(dst->attrs(),
                                     outside_compilation_attr_name,

  TF_OperationDescription* init(string type) {
    // Construct op_name to match the name used by REGISTER_OP in the
    // ATTR_TEST_REGISTER calls above.
    string op_name = "CApiAttributesTestOp";
    if (type.find("list(") == 0) {
      op_name += "List";
      type = type.replace(0, 5, "");
      type = type.replace(type.size() - 1, 1, "");
    }
    op_name += type;
    return TF_NewOperation(

    LLVM_DEBUG(llvm::dbgs() << "Zero point value should be a tensor type. Got: "
                            << zero_points_type << ".\n");
    return failure();
  }

  if (Type zero_points_element_type = zero_points_type.getElementType();
      !mlir::isa<IntegerType>(zero_points_element_type)) {
    LLVM_DEBUG(llvm::dbgs() << "Zero point should be an integer type. Got: "
                            << zero_points_element_type << ".\n");

  // First create a simple identity function.
  TF_Graph* function_graph = TF_NewGraph();
  TF_OperationDescription* arg_descr =
      TF_NewOperation(function_graph, "Placeholder", "arg");
  TF_SetAttrType(arg_descr, "dtype", TF_INT32);
  TF_Status* status = TF_NewStatus();
  TF_Operation* arg = TF_FinishOperation(arg_descr, status);
  ASSERT_TRUE(TF_GetCode(status) == TF_OK) << TF_Message(status);
  TF_OperationDescription* id_descr =

        llvm::dyn_cast<TF::ConstOp>(fold_results[0].getDefiningOp()).getValue(),
        &tensor));
    xla::Shape xla_shape;
    TF_RETURN_IF_ERROR(tensorflow::TensorShapeToXLAShape(
        tensor.dtype(), tensor.shape(), &xla_shape));
    xla::PjRtClient::HostBufferSemantics host_buffer_semantics =
        xla::PjRtClient::HostBufferSemantics::kImmutableUntilTransferCompletes;
    TF_ASSIGN_OR_RETURN(
        auto buffer,

}

Node* KnownShapeBase(DataType dtype, absl::Span<const int> shape,
                     const GraphDefBuilder::Options& opts) {
  if (opts.HaveError()) return nullptr;
  NodeBuilder node_builder(opts.GetNameForOp("Const"), "Const",
                           opts.op_registry());
  TensorProto value;
  value.set_dtype(dtype);
  for (int dim : shape) {