// RUN: tf-opt %s -tfl-partitioned-topological-sort | FileCheck %s

// CHECK-LABEL: @tf_ops_will_be_partitioned
func.func @tf_ops_will_be_partitioned() -> tensor<1xf32> {
  %const = "tfl.pseudo_const"() {value = dense<[1.0]> : tensor<1xf32>} : () -> tensor<1xf32>
  %tmp1 = "tfl.add"(%const, %const) { fused_activation_function = "NONE" } : (tensor<1xf32>,tensor<1xf32>) -> (tensor<1xf32>)

// 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
  }

    }
    # This should be modeled as unit attribute
    attr {
      key: "_input_shapes"
      value {
      }
    }
  }
}
versions {
  producer: 85
  min_consumer: 12
}

# Match partitioned call in main and capture the callee name.
# CHECK-LABEL: func @main
# CHECK: tf.PartitionedCall
# CHECK-SAME: f = @[[FUNCTION:[a-zA-Z0-9_]*]]

# Verify that callee has the unit attribute tf._input_shapes.

  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}> : () -> ()