std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> matmul(
      MatMulOp(context, hcpu, hdevice), TFE_DeleteOp);
  TFE_OpSetDevice(matmul.get(), name, status.get());
  ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get());
  TFE_TensorHandle* retval;
  int num_retvals = 1;
  TFE_Execute(matmul.get(), &retval, &num_retvals, status.get());

category names correspond to generated source file names, and should be
consistent with the original source files registering each operator. For example
since `REGISTER_OP("MatMul")` appears in ***core/math_ops.cc***, the "MatMul"
operator in the script should be in the "math" category, and it will be
generated in the output file `c/experimental/ops/math_ops.cc`.

    // Calc Grad
    AbstractTensorHandle* matmul_A_output;
    AbstractTensorHandle* matmul_B_output;
    std::string name_grad_A = "MatMul_Grad_A";
    std::string name_grad_B = "MatMul_Grad_B";
    if (!t_a && !t_b) {
      TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, B.get(), &matmul_A_output,
                                /*transpose_a = */ false,

  }

  TFE_Op* matmul = nullptr;
  if (func) {
    const string matmul_device = remote_func_outputs ? task2_name : "";
    string function_def = MatMulFunction(matmul_device);
    TFE_ContextAddFunctionDef(ctx, function_def.data(), function_def.size(),
                              status);
    CHECK_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);

    matmul = TFE_NewOp(ctx, "MatMulFunction", status);

  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  TFE_Op* matmul = MatMulOp(ctx, h0_task1, h1_task1);
  TFE_OpSetDevice(matmul, remote_device_name, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  TFE_TensorHandle* retvals[1];
  int num_retvals = 1;
  TFE_Execute(matmul, &retvals[0], &num_retvals, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  TFE_TensorHandle* m = TestMatrixTensorHandle(ctx);
  TFE_Op* matmul = TFE_NewOp(ctx, "MatMul", status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_TensorHandle* retvals[1];
  int num_retvals = 1;
  for (auto s : state) {
    TFE_OpReset(matmul, "MatMul", nullptr, status);
    CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
    TFE_OpAddInput(matmul, m, status);

              absl::Span<AbstractTensorHandle* const> inputs,
              absl::Span<AbstractTensorHandle*> outputs) -> Status {
        return ops::MatMul(ctx, inputs[0], inputs[1], &outputs[0], transpose_a,
                           transpose_b, "MatMul");
      };
      ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
          MatMulModel, BuildGradModel(MatMulModel, registry_),

  // Build an abstract operation.
  auto* matmul_op = TF_NewAbstractOp(graph_ctx);
  TF_AbstractOpSetOpType(matmul_op, "MatMul", status.get());
  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());
  TF_AbstractOpSetOpName(matmul_op, "my_matmul", status.get());
  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());

  TFE_Context* tfe_context_;
};

TEST_F(ShapeInferenceTest, InfersShapesFromInputShapes) {
  TFE_Op* matmul_op;
  matmul_op = TFE_NewOp(tfe_context_, "MatMul", status_);
  CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);

  // Infer shape when everything is known.
  CheckOutputShapes(matmul_op,
                    /*input_shapes*/ {make_shape({3, 2}), make_shape({2, 4})},
                    /*input_tensors*/ {},

  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  TFE_Op* matmul = MatMulOp(ctx, h0_task1, h1_task1);
  TFE_OpSetDevice(matmul, remote_device_name, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  TFE_TensorHandle* retvals[1];
  int num_retvals = 1;
  TFE_Execute(matmul, &retvals[0], &num_retvals, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);