auto launch = builder->create<tf_device::LaunchOp>(
      before_op->getLoc(), builder->getStringAttr(host_device),
      launch_result_types);
  launch.getBody().push_back(launch_block);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<tf_device::ReturnOp>(before_op->getLoc(), launch_results);

  return launch;
}

// Creates a new_launch after `before_op` with `ops_to_move` from

  // supported.
  const auto& device_list =
      devices.find(tensorflow::GetDeviceAliasForLogicalCore(0))->getSecond();

  // Create the state variable for each device.
  for (llvm::StringRef device : device_list) {
    state_vars.push_back(builder->create<TF::VarHandleOp>(
        loc,
        llvm::ArrayRef<Type>{RankedTensorType::get(
            {}, TF::ResourceType::get(llvm::ArrayRef<TensorType>{key_type},

    auto static_shape_attr =
        mlir::DenseIntElementsAttr::get(static_shape_type, static_shape);
    return rewriter.create<TF::ConstOp>(loc, static_shape_attr).getOutput();
  }

  // If the shape is not static, create a new ShapeOp.
  BoolAttr false_attr = rewriter.getBoolAttr(false);
  return rewriter
      .create<TF::ShapeOp>(loc, input,
                           /*use_32bit=*/false_attr)
      .getOutput();
}

        return false;
      }
    }
    rewriter.setInsertionPointAfter(op);
    auto q = rewriter.create<quantfork::QuantizeCastOp>(op->getLoc(), cast_type,
                                                        op.getResult());
    auto dq = rewriter.create<quantfork::DequantizeCastOp>(op->getLoc(),
                                                           expressed_type, q);

    TF_RETURN_IF_ERROR(c_event->Create());
    return std::move(c_event);
  }

  absl::StatusOr<std::unique_ptr<Stream>> CreateStream(
      std::optional<std::variant<StreamPriority, int>> priority =
          std::nullopt) override {
    auto stream = std::make_unique<CStream>(&device_, stream_executor_, this);
    TF_RETURN_IF_ERROR(stream->Create());
    return std::move(stream);
  }

 private:

  EXPECT_TRUE(xla::LiteralTestUtil::Equal(
      *literal1, xla::LiteralUtil::CreateR2<int32_t>({{1, 2, 3}})));

  std::shared_ptr<xla::Literal> literal2 = *exec_args[1]->ToLiteralSync();
  EXPECT_TRUE(xla::LiteralTestUtil::Equal(
      *literal2, xla::LiteralUtil::CreateR2<int32_t>({{4, 5, 6}})));
}

TEST_F(PjRtExecutionUtilTest, PreparePjRtExecutableArgumentsVariableInputs) {

    auto new_shape_const = rewriter.create<arith::ConstantOp>(
        weight_op->getLoc(), shape_spec_type, new_shape_const_attr);
    auto reshape_op = rewriter.create<TF::ReshapeOp>(
        weight_op->getLoc(), new_shape, weight_op->getResult(0),
        new_shape_const);
    op->setOperand(weight_operand_idx, reshape_op);

    // Create a new function with preprocessed types.

    rewriter.setInsertionPointAfter(preceding_sc_op);
    auto fcast_op = rewriter.create<TF::CastOp>(
        preceding_sc_op->getLoc(), dq_arg_type.clone(rewriter.getF32Type()),
        preceding_sc_op.getResult());

    // Create a new AvgPool op with float type.
    TF::AvgPoolOp float_avg_pool_op = rewriter.create<TF::AvgPoolOp>(
        avg_pool_op->getLoc(),
        avg_pool_op.getType().clone(rewriter.getF32Type()),

    IntegerAttr zp_attr;
    if (!matchPattern(zp, m_Constant(&zp_attr))) {
      return failure();
    }
    rewriter.setInsertionPoint(zp.getDefiningOp());
    auto zp_tensor = rewriter.create<TF::ConstOp>(
        loc, RankedTensorType::get({}, zp.getType()), zp_attr);
    auto zp_cast = rewriter.create<CastOp>(

  context.loadDialect<mlir::tf_device::TensorFlowDeviceDialect>();
  mlir::OwningOpRef<mlir::ModuleOp> module_ref =
      mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
  mlir::OpBuilder builder(module_ref->getBodyRegion());

  llvm::SmallVector<mlir::Type, 8> result_types;
  auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      mlir::UnknownLoc::get(&context), result_types);
  cluster->setAttr(kNumCoresPerReplicaAttr,