%cst = "tf.Const"() {device = "", value = dense<[[[[-0.315365672, 0.27481091], [0.0901821703, -0.382271349], [-0.105572946, -0.354302853]], [[-0.47703138, -0.307006568],...

module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 29 : i32}} {
  func.func @main() {
    tf_executor.graph {
      %0 = tf_executor.island wraps "tf.PartitionedCall"() {Tin = [], Tout = [], config = "", config_proto = "", device = "", executor_type = "", f = @foo, name = "Call_foo"} : () -> ()
      tf_executor.fetch
    }
    func.return
  }
  func.func @foo() {

    %1 = "tf.Reshape"(%arg1, %0) : (tensor<19x10xf32>, tensor<2xi64>) -> tensor<?x?xf32>

    // CHECK: %[[RESULT:.*]] = mhlo.constant dense<[10, 19]>
    %2 = "tf.PartitionedCall"(%1) {config = "", config_proto = "", executor_type = "", f = @get_shape} : (tensor<?x?xf32>) -> (tensor<?xi64>)

    // CHECK: return %[[RESULT]]
    func.return %2 : tensor<?xi64>
  }

  // CHECK: _xla_original_oc_node_name = [[NODE_NAME1:.*]], _xla_token_input_nodes = ["_xla_token_arg_node"]
  %0 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @PartitionedCall3} : (tensor<i32>) -> (tensor<i32>)
  // CHECK-NOT: _xla_token_input_nodes
  %1 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @IdentityFunc} : (tensor<i32>) -> (tensor<i32>)

func.func @serving_default(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x2x2x2xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input_tensor:0", outputs = "PartitionedCall:0"}} {

    func.return %2 : tensor<*xi32>
  }
}

// -----

// CHECK: module attributes {tfl._legalize_tfl_variables = true}
module {
  func.func @main() -> tensor<*xi32> {
    %0 = "tf.PartitionedCall"() {f = @f, config = "", config_proto = "", executor_type = ""}
      : () -> tensor<*xi32>
    func.return %0 : tensor<*xi32>
  }
  func.func @f() -> tensor<*xi32> {

  Option<bool> skip_partitioned_calls_{
      *this, "skip-partitioned-calls",
      ::llvm::cl::desc(
          "Skip tf.StatefulPartitionedCall and tf.PartitionedCall")};
};

void TFToMhloPass::runOnOperation() {
  auto func = getOperation();
  MLIRContext *context = func->getContext();

  RewritePatternSet patterns(context);

  auto grad = SymbolicGradient(scope, func_inputs, input_dtypes, f);
  for (int i = 0; i < num_inputs; i++) {
    grad_outputs->push_back(grad[i]);
  }

  return scope.status();
}

REGISTER_GRADIENT_OP("PartitionedCall", PartitionedCallGrad);
REGISTER_GRADIENT_OP("StatefulPartitionedCall", PartitionedCallGrad);

}  // anonymous namespace
}  // namespace ops

  return %arg0 : tensor<f32>
}

// CHECK-LABEL: handles_partitioned_function_calls
func.func @handles_partitioned_function_calls(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
  // CHECK: PartitionedCall"()
  %1 = "tf.PartitionedCall"(%arg0, %arg1) {f = @f} : (tensor<f32>, tensor<f32>) -> tensor<f32>
  return %1 : tensor<f32>
}

// -----

func.func private @f(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {

    %2:2 = tf_executor.island wraps "tf.StatefulPartitionedCall"(%0#0, %arg1) {Tin = ["tfdtype$DT_FLOAT", "tfdtype$DT_RESOURCE"], Tout = ["tfdtype$DT_FLOAT"], _gradient_op_type = "PartitionedCall-1205", config = "", config_proto = "\0A\07\0A\03GPU\10\00\0A\07\0A\03CPU\10\012\02J\008\01", device = "", executor_type = "", f = @function0} : (tensor<f32>, tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>>) -> tensor<f32> loc("statefulpartitionedcall")...