// function.
  SmallVector<TF::XlaCallModuleOp> xla_call_module_ops;
  for (auto xla_call_module_op : func_ops[0].getOps<TF::XlaCallModuleOp>()) {
    xla_call_module_ops.push_back(xla_call_module_op);
  }
  ASSERT_THAT(xla_call_module_ops, SizeIs(1));
  auto xla_call_module_op = xla_call_module_ops[0];
  EXPECT_THAT(xla_call_module_op,

  if (failed(solver.initializeAndRun(module))) return signalPassFailure();

  DenseSet<GlobalTensorOp> remaining_global_tensor_ops;
  {
    auto ops = module.getOps<GlobalTensorOp>();
    remaining_global_tensor_ops.insert(ops.begin(), ops.end());
  }

  for (auto global_tensor : remaining_global_tensor_ops) {
    // This pass assumes that all global tensors as immutable (e.g. by a

  TF_Tensor* analytical_tensor;
  auto s = GetValue(t, &analytical_tensor);
  ASSERT_EQ(errors::OK, s.code()) << s.message();

  int64_t num_elem_analytical = 1;
  auto num_dims_analytical = TF_NumDims(analytical_tensor);
  ASSERT_EQ(dims.size(), num_dims_analytical);
  for (int j = 0; j < num_dims_analytical; j++) {
    auto dim_analytical = TF_Dim(analytical_tensor, j);

namespace mlir {
namespace TF {

using ::tensorflow::kValidDeviceTypes;

LogicalResult HasValidCompilationAndReplicationAttributes(Operation& op) {
  auto replicate_attr = op.getAttrOfType<StringAttr>(kReplicationInfoAttr);
  auto compile_attr = op.getAttrOfType<StringAttr>(kCompileDeviceTypeAttr);
  if (!replicate_attr && !compile_attr) return success();
  if (!replicate_attr || !compile_attr)

      platform->ExecutorForDevice(0).value();
  TF_ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());

  se::DeviceMemoryBase gpu_dst{device_tensor.data(), 4 * sizeof(float)};
  xla::Shape shape;
  TF_ASSERT_OK(TensorShapeToXLAShape(DT_FLOAT, TensorShape({2, 2}), &shape));

  TF_ASSERT_OK_AND_ASSIGN(auto se_event, executor->CreateEvent());
  tsl::AsyncValueRef<std::unique_ptr<se::Event>> done_event =

    auto element_type = attr.getType().getElementType();
    if (attr.getNumElements() != 1 || !element_type.isIntOrFloat())
      return false;
    if (mlir::isa<FloatType>(element_type)) {
      auto value = *attr.value_begin<APFloat>();
      return value.isNegative() && value.isInfinity();
    } else if (element_type.isInteger(1)) {
      auto value = *attr.value_begin<APInt>();

    auto transpose_op = reduce_input.getDefiningOp<stablehlo::TransposeOp>();
    if (!transpose_op || !transpose_op->hasOneUse()) return failure();
    Value transpose_input = transpose_op.getOperand();

    ArrayRef<int64_t> transpose_perm = transpose_op.getPermutation();

    // Inversely transposes all the attributes to prepare for the new reduce op
    auto getInvertPermutedAttr =

  llvm::SmallVector<Type, 4> operand_types;
  operand_types.reserve(live_ins.size());
  for (Value v : live_ins) operand_types.emplace_back(v.getType());

  auto func_type = builder->getFunctionType(operand_types, op.getResultTypes());

  std::string func_name;
  if (auto outlined_func_name = op->template getAttrOfType<StringAttr>(
          TF::kClusterOutlinedFunctionNameAttr)) {
    op->removeAttr(TF::kClusterOutlinedFunctionNameAttr);

    // We may retrieve the original function's name from the attribute.
    // Functions without this attribute are ignored.
    auto original_func_name =
        func_op->getAttrOfType<StringAttr>("tf._original_func_name");
    if (original_func_name) {
      if (auto alias_itr = function_aliases.find(original_func_name.str());
          alias_itr != function_aliases.end()) {

inline bool TFDataFormatIsNHWC(Operation *op) {
  auto attr = op->getAttrOfType<StringAttr>("data_format");
  return !attr || attr.getValue() == "NHWC";
}

// Returns true if the given TensorFlow op does not have a `data_format`
// attribute (then default to "NDHWC"), or its `data_format` attribute is
// "NDHWC".
inline bool TFDataFormatIsNDHWC(Operation *op) {
  auto attr = op->getAttrOfType<StringAttr>("data_format");