// RUN: odml-to-stablehlo-opt %s --tf-stablehlo=skip-partitioned-calls=true | FileCheck %s --check-prefix=CHECK-SKIP
// RUN: odml-to-stablehlo-opt %s --tf-stablehlo=skip-partitioned-calls=false | FileCheck %s --check-prefix=CHECK-NOSKIP

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

The new [Cookie.Partitioned] field identifies cookies with the Partitioned attribute....

//
// The class currently only interacts with single-slice (i.e., non-partitioned)
// variables.
class CheckpointReader {
 public:
  CheckpointReader(const string& filename, TF_Status* status);

  bool HasTensor(const string& name) const;
  const string DebugString() const;

  // Returns a map from variable names to their shapes.  Slices of a partitioned
  // tensor are combined into a single entry.

func.func @stateful_pcall_func(%arg0: tensor<i32>) -> tensor<i32> {
  func.return %arg0 : tensor<i32>
}

// -----

// Check that we only outline the outermost partitioned call with
// `_xla_compile_device_type` in nested calls to a device cluster. The callee
// of any partitioned call outside of a device cluster is marked with
// `tf._noinline = true` to prevent it getting inlined to perseve runtime
// performance.

using mlir::SymbolUserOpInterface;
using mlir::func::FuncOp;

#define GEN_PASS_DEF_XLACLUSTERFORMATIONPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"

// Outlines partitioned call ops with `_XlaMustCompile` to device clusters.
struct XlaClusterFormationPass
    : public impl::XlaClusterFormationPassBase<XlaClusterFormationPass> {
  void runOnOperation() override;
};

        "tf.StatefulPartitionedCall"(%arg0) <{config = "", config_proto = "", executor_type = "", f = @partitioned}> : (tensor<*x!tf_type.resource>) -> ()
        tf_executor.yield
      }
      tf_executor.fetch %control : !tf_executor.control
    }
    return
  }

  Option<bool> skip_resize_{
      *this, "skip-resize",
      ::llvm::cl::desc("Skip tf.ResizeBilinear and tf.ResizeNearestNeighbor")};

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

void TFToMhloPass::runOnOperation() {
  auto func = getOperation();

  func.func private @f() {
    func.call @g() : () -> ()
    func.return
  }
  func.func private @g() {
    func.call @f() : () -> ()
    func.return
  }
}

// -----
// Test stateful and stateless partitioned calls.
// CHECK-LABEL: func @f
func.func @f() {
  // CHECK: "tf.PartitionedCall"() <{config = "",  config_proto = "", executor_type = "", f = @g}> : () -> ()

	})
}

var (
	requestTotal = compbasemetrics.NewCounterVec(
		&compbasemetrics.CounterOpts{
			Name:           "apiserver_delegated_authz_request_total",
			Help:           "Number of HTTP requests partitioned by status code.",
			StabilityLevel: compbasemetrics.ALPHA,
		},
		[]string{"code"},
	)

	requestLatency = compbasemetrics.NewHistogramVec(
		&compbasemetrics.HistogramOpts{

//   Executor Dialect or TF Functional Dialect. Will convert to TF Functional.
//   is_supported_by_replicated_brige - If the graph targets the replicated
//   bridge. Set it to true for replicated/partitioned graphs. e.g. replicated
//   and single-core TPU graphs. Set this to false if the graph is not
//   replicated, e.g. CPU/GPU graphs. is_in_fallback_enabled_mode - Whether this