// CHECK: %[[GVAR2:.*]] = "tf.MlirLocalVarOp"() : () -> tensor<!tf_type.resource<tensor<5x3xf32>>>
  // CHECK: "tf.StatefulPartitionedCall"(%[[VAR]], %[[GVAR1]], %[[GVAR2]])
  // CHECK-SAME: f = @callee_tensorarray_decomposed
  %call = "tf.StatefulPartitionedCall"(%ta#0) {f = @callee, config = "", config_proto = "", executor_type = ""}
    : (tensor<!tf_type.resource>) -> tensor<!tf_type.resource>

  // CHECK: %[[CLUSTER:.*]] = "tf_device.cluster"()
  "tf_device.cluster"() ({
    // CHECK: %[[PC0:.*]] = "tf.StatefulPartitionedCall"(%[[READ0]], %[[READ1]], %[[CONST]])
    // CHECK-SAME: f = @callee_resource_lifted
    %3 = "tf.StatefulPartitionedCall"(%0, %1, %2) {f = @callee, config = "", config_proto = "", executor_type = ""}

  // 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 = [{