// when the producer is a non-XLA TF GPU kernel or function (e.g.
    // tf.matmul).
    //
    // 3. AsyncValueTensor, containing a PjRtBuffer. This is the legacy mode
    // and certain device type (e.g. TPU) still uses this path.
    AsyncValueTensor* av_tensor = AsyncValueTensor::FromTensor(tensor);
    if (use_pjrt_tensor_buffer) {
      if (av_tensor != nullptr) {
        return absl::InvalidArgumentError(

    return stage == IrExportStage::HLO ||
           stage == IrExportStage::HLO_NO_METADATA ||
           stage == IrExportStage::HLO_SERIALIZED;
  };
  // TODO(b/238830423): support GetCompilerIr on TFRT TPU device for stages
  // that requires compilation from HLO to executable.
  if (device_type != DEVICE_CPU && stream == nullptr &&
      !is_tfrt_tpu_supported_stage(stage)) {
    return absl::InternalError(

// call tree upwards and duplicates all constants found in the subtree in a
// similar manner.
//
// This pass may be used to avoid placing shape-determining constants in the CPU
// graph and pass them as arguments to the TPU graph (via `TPUPartitionedCall`).
// If this happens, the XLA compiler cannot recognize such arguments as
// constants and may result in an error.
//
// A set of predefined ops and operand indices is used to determine whether an

// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-tpu-resource-partition | FileCheck %s

func.func private @computation(%arg0: tensor<i32>) -> tensor<i32>

// CHECK-LABEL: func @read_write_resource
// CHECK-SAME: ([[ARG0:%.+]]: tensor<!tf_type.resource<tensor<i32>>>, [[ARG1:%.+]]: tensor<!tf_type.resource<tensor<i32>>>)
func.func @read_write_resource(%arg0: tensor<!tf_type.resource<tensor<i32>>>, %arg1: tensor<!tf_type.resource<tensor<i32>>>) {

#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/core_platform_payloads.pb.h"
#include "tensorflow/core/tpu/tpu_defs.h"
#include "tensorflow/core/util/debug_data_dumper.h"
#include "tsl/platform/errors.h"

namespace tensorflow {
namespace {
constexpr absl::string_view kGroupSizeAttrName =

  func.return  %1 : tensor<*xf32>

}

}

// -----

module attributes {tf_saved_model.semantics} {

// Test not hoisting in TPU functions.

// CHECK-LABEL: func @_tfrt_resource_init
// CHECK-NEXT: return

// CHECK-LABEL: func private @func2
func.func private @func2(%arg: tensor<i1>) -> tensor<i32> {
  // CHECK-NOT: tf._TfrtGetResource

// Default version number for native serialization.
constexpr int64_t kDefaultVersion = 9;
// Platforms for XlaCallModuleOp.
constexpr StringRef kPlatformCpu = "CPU";
constexpr StringRef kPlatformTpu = "TPU";

class ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass
    : public impl::
          ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPassBase<

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

#define DEBUG_TYPE "tf-tpu-merge-variables-with-execute"

namespace mlir {
namespace TFTPU {

namespace {
constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr char kDeviceAttr[] = "device";

  // Current supported ops between convolution input and the block arguments are
  // PadOp and CastOp.
  return GetInputBlockArgNum(conv2d.getInput());
}

// Applies space to depth transform for the first convolution on TPU device.
void HandleFirstConvolution(TF::Conv2DOp conv2d, int64_t block_size) {
  // Check if input and filter type are RankedTensorType.
  auto input_tensor_type =

    // handles empty computations. Another approach could be to drop empty
    // clusters within MLIR but that seems to trigger other failures but can be
    // considered again.
    // Old bridge only removes unsupported TPU types (only string for now)
    // during outside compilation extraction so this should be enough for
    // the parity.
    bool is_unsupported_type = mlir::isa<mlir::TF::StringType>(