const std::vector<Output>& grad_inputs,
                        std::vector<Output>* grad_outputs) {
  ops::GradFunc grad_fn;
  TF_RETURN_IF_ERROR(ops::GradOpRegistry::Global()->Lookup(
      op.node()->type_string(), &grad_fn));
  TF_RETURN_IF_ERROR(grad_fn(scope, op, grad_inputs, grad_outputs));
  TF_RETURN_IF_ERROR(scope.status());
  return absl::OkStatus();
}

}  // end namespace test

  ops::GradFunc grad_fn;
  Status s = registry_->Lookup(opname, &grad_fn);
  return s.ok() && (grad_fn == nullptr);
}

Status SymbolicGradientBuilder::CallGradFunction(
    const Operation& op, const std::vector<Output>& grad_inputs,
    std::vector<Output>* grad_outputs) {
  ops::GradFunc grad_fn;
  TF_RETURN_IF_ERROR(registry_->Lookup(op.node()->type_string(), &grad_fn));

    reserve_spaces[1] = op.input(4);  // pop_var
    if (data_format == "NCHW") {
      x = Transpose(scope, x, {0, 2, 3, 1});
      grad_y = Transpose(scope, grad_y, {0, 2, 3, 1});
    } else if (data_format == "NCDHW") {
      x = Transpose(scope, x, {0, 2, 3, 4, 1});
      grad_y = Transpose(scope, grad_y, {0, 2, 3, 4, 1});
    }

    StringPiece target_data_format;
    if (data_format == "NCHW" || data_format == "NHWC") {

def ConvertTFMatMulToXLADotV2Op : Pat<
  (TF_MatMulOp:$matmul
    (TF_SubOp (TF_CastOp $input, $truncate), $input_zp),
    (TF_CastOp (TF_IdentityOp $weight), $truncate1),
    $transpose_a, $transpose_b, $grad_a, $grad_b),
  (CreateXlaDotV2OpFromTfMatMulOp
    $input, $weight, $input_zp,
    /*weight_zp=*/(CreateScalarIntegerConst<"int32_t", "0">), $matmul,
    $transpose_a, $transpose_b),
  [(IsInt8ElementType $input),

      (NamedAttr<"dilations"> $dilations))),
  [(IsNotInLiftedFunc $res)], [], (addBenefit 1)>;

def LiftMatMul : Pat<
  (TF_MatMulOp:$res $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
  (LiftAsTFPartitionedCall<"composite_matmul_fn">
    (ArgumentList $a, $b),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"transpose_a"> $transpose_a),

      (NamedAttr<"dilations"> $dilations))),
  [(IsNotInLiftedFunc $res), (IsConstTensor $filter)], [], (addBenefit 1)>;

def LiftMatMul : Pat<
  (TF_MatMulOp:$res $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
  (LiftAsTFPartitionedCall<"composite_matmul_fn">
    (ArgumentList $a, $b),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"transpose_a"> $transpose_a),

  // CHECK: %[[MATMUL_4:.*]] = "tf.MatMul"(%[[LHS_4]], %[[RHS_4]]) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> : (tensor<4x5xf32>, tensor<5x6xf32>) -> tensor<4x6xf32>

  }

  func.func private @composite_matmul_with_bias_fn_1(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>, %arg2: tensor<3xf32>) -> tensor<1x3xf32> attributes {tf_quant.composite_function} {
    %0 = "tf.MatMul"(%arg0, %arg1) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> {attr_map = "0:transpose_a,1:transpose_b", device = ""} : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>

    $adj_x, $adj_y, $grad_x, $grad_y),
  (TF_BatchMatMulV2Op:$dest $x, $y, $adj_x, $adj_y, $grad_x, $grad_y),
  [], [(CopyAttrs $src, $dest)]>;

def BatchMatMulToMatMul : Pat<
  (TF_BatchMatMulOp:$src $x, $y, $adj_x, $adj_y, $grad_x, $grad_y),
  (TF_MatMulOp:$dest $x, $y, $adj_x, $adj_y, $grad_x, $grad_y),
  [(IsRank2Tensor $x), (IsRank2Tensor $y)],
  [(CopyAttrs $src, $dest)]>;

  TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("transpose_a", transpose_a));
  TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("transpose_b", transpose_b));
  TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("grad_a", grad_a));
  TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("grad_b", grad_b));
  int num_retvals = 1;
  return op_ptr->Execute(absl::MakeSpan(product, 1), &num_retvals);
}

// Op: IdentityN()
// Summary:
//
// Description: