// -----

// This test makes sure functions without tf._noinline=true is inlined.

module {
  func.func @partitioned_call(%arg0: tensor<1x2x2x3xf32>) -> (tensor<1x2x2x3xf32>) {
    %0 = "tf.PartitionedCall"(%arg0) <{
      config = "", config_proto = "", executor_type = "", f = @some_func
    }> {
      _collective_manager_ids = [], device = ""
    } : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>

  Output in = ops::Placeholder(root, DT_FLOAT);
  Output perm = ops::Const(root, {3, 1, 2, 0});

  NameAttrList b_name_attr;
  b_name_attr.set_name("TransposeCall");
  ops::PartitionedCall call(root.WithOpName("call"), {in, perm}, {DT_FLOAT},
                            b_name_attr);
  call.output.front().node()->AddAttr(kXlaMustCompileAttr, true);

  std::unique_ptr<Graph> graph;

  }
  attr {
    key: "shape"
    value {
      shape {
      }
    }
  }
  attr {
    key: "shared_name"
    value {
      s: ""
    }
  }
}
node {
  name: "Call"
  op: "PartitionedCall"
  attr {
    key: "Tin"
    value {
      list {
      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_RESOURCE
      }
    }
  }
  attr {

  auto value = device_attr.getValue();
  // TODO(b/229028654): Remove string conversion once we have C++17.
  absl::string_view device_type(value.data(), value.size());
  // Device type may be empty for some ops, e.g. tf.PartitionedCall.
  auto it = std::find(kValidDeviceTypes.begin(), kValidDeviceTypes.end(),
                      device_type);
  if (it == kValidDeviceTypes.end()) return failure();
  return success();
}

      }
    }
  }
  attr {
    key: "Tout"
    value {
      list {
        type: DT_FLOAT
      }
    }
  }
  attr {
    key: "_gradient_op_type"
    value {
      s: "PartitionedCall-1205"
    }
  }
  attr {
    key: "config"
    value {
      s: ""
    }
  }
  attr {
    key: "config_proto"
    value {

  // CHECK-NOT: Optional
  %0 = "tf.OptionalFromValue"(%arg0) : (tensor<f32>) -> tensor<!tf_type.variant<tensor<f32>>>
  %1 = "tf.PartitionedCall"(%0) {
      config = "",
      config_proto = "",
      executor_type = "",
      f = @identity} : (tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>>

    %0 = tf_executor.graph {
      %out, %ctl = tf_executor.island wraps "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @serving_default} : (tensor<?x!tf_type.string>) -> tensor<*xi64>
      tf_executor.fetch %out : tensor<*xi64>
    }
    return %0 : tensor<*xi64>
  }

      return failure();
    }
    rewriter.setInsertionPointAfter(call_op);

    // The TPUPartitionedCall has a TPUOrdinalSelectorOp for its last argument
    // which should be removed. So the replaced PartitionedCall op should keep
    // its original arguments except for the last element.
    SmallVector<Value> args = call_op.getOperands().drop_back();

    rewriter.replaceOpWithNewOp<TF::PartitionedCallOp>(

    .SetShapeFn([](InferenceContext* c) {
      c->set_output(0, c->input(0));
      return absl::OkStatus();
    })
    .Doc(R"(XLA Merge Op. For use by the XLA JIT only.

Merges the outputs from the PartitionedCall node and the _XlaRun node.
Unlike the TensorFlow Merge op, which requires inputs of some types to be
placed on the host, the _XlaMerge op can merge inputs of all types when

        tf_executor.graph {
          %control_0 = tf_executor.island wraps "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @function_1}> : () -> ()
          %control_1 = tf_executor.island wraps "tf.PartitionedCall"() <{config = "", config_proto = "", executor_type = "", f = @function_2}> : () -> ()