// TPUCompilationResultsOp.
  if (op->hasAttr(kTpuCompilationStatus)) {
    op->setAttr(kTpuCompilationStatus, attr);
  }
}

// A helper class to inline TF::StatefulPartitionedCall ops
struct Inliner : public InlinerInterface {
  Inliner(OpBuilder& builder, SymbolTable& symbol_table)
      : InlinerInterface(builder.getContext()),
        builder(builder),

      "SpaceToBatchND",
      "SpaceToDepth",
      "SparseMatMul",
      "SparseToDense",
      "StackCloseV2",
      "StackPopV2",
      "StackPushV2",
      "StackV2",
      "StatefulPartitionedCall",
      "StatefulStandardNormalV2",
      "StatefulTruncatedNormal",
      "StatefulUniform",
      "StatefulUniformFullInt",
      "StatefulUniformInt",
      "StatelessCase",

    TF_Tensor:$output
  );

  TF_DerivedResultTypeAttr dtype = TF_DerivedResultTypeAttr<0>;
  DerivedAttr shape = TF_DerivedResultShapeAttr;
}

def TF_StatefulPartitionedCallOp : TF_Op<"StatefulPartitionedCall",
      [CallOpInterface, DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
  let summary =
    "returns `f(inputs)`, where `f`'s body is placed and partitioned.";

  let description = [{