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

				if err != nil {
					c.Expires = time.Time{}
					break
				}
			}
			c.Expires = exptime.UTC()
			continue
		case "path":
			c.Path = val
			continue
		case "partitioned":
			c.Partitioned = true
			continue
		}
		c.Unparsed = append(c.Unparsed, parts[i])
	}
	return c, nil
}

// readSetCookies parses all "Set-Cookie" values from

}
}  // namespace

// Rewrites unpartitioned resource reads and writes to partitioned resource
// reads and writes. The TPU computation from the frontend is generated in such
// a way that resource operations operate on the unpartitioned resource handle
// (from a `tf.TPUReplicatedInput`). This results in resource reads and writes
// on the unpartitioned resource handle post resource op decomposition/lifting.

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

using PartitionedClusterOutputMap = absl::flat_hash_map<
    uint64_t, llvm::SmallVector<mlir::TF::TPUPartitionedOutputV2Op, 8>>;

// Returns the partitioned output ops from the cluster if there are any,
// along with any single user identity ops between them. Not all outputs
// of a cluster must be partitioned, so the output is a map from cluster
// output ids to ops.
std::tuple<PartitionedClusterOutputMap,