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.

namespace {

// Helper class for parsing operands to a foldable operation.
class FoldAdaptor {
 public:
  // Returns std::nullopt if the operation cannot be folded.
  static std::optional<FoldAdaptor> Create(Operation* operation) {
    auto foldable_opr = [](Value val) -> bool {
      return !llvm::isa<BlockArgument>(val) &&
             llvm::isa<stablehlo::ConstantOp>(val.getDefiningOp());
    };

  TensorShape shape(shape_vector);

  // Create the variable.
  Status status;
  std::unique_ptr<Variable> var;
  TF_EXPECT_OK(Variable::CreateUninitialized(context(), dtype, shape,
                                             absl::nullopt, nullptr, {}, &var));

  // Create a TensorHandle
  ImmediateTensorHandlePtr expected_handle =

std::optional<Value> MaterializeIllegalCast(OpBuilder &builder, Type type,
                                            ValueRange inputs, Location loc) {
  return builder.create<UnrealizedConversionCastOp>(loc, type, inputs)
      ->getResult(0);
}

class StablehloToOdmlTypeConverter : public vhlo::VhloTypeConverter {
 public:
  StablehloToOdmlTypeConverter() : vhlo::VhloTypeConverter() {

    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()),

  *library_expected.add_function() = FunctionDefHelper::Create(
      "F1", {"a_0_arg:float"}, {"c_0_retval:float"}, {},
      {
          {{"C"}, "UnaryTest", {"a_0_arg"}},
      },
      {{"c_0_retval", "C:o:0"}});
  *library_expected.add_function() = FunctionDefHelper::Create(
      "F2", {"b_0_arg:float", "c_0_arg:float"}, {"d_0_retval:float"}, {},
      {

    bool used_persistent_cache) {
  metrics::UpdateXlaCompilationTime(compile_time_us);

  const std::string& function_name = function.name();

  mutex_lock lock(mu_);
  // Create a stats entry if it doesn't already exist.
  auto it =
      cluster_compile_stats_.emplace(function.name(), ClusterCompileStats{})
          .first;

  const uint64 compile_time_s = compile_time_us / 1.0e6;

  core::RefCountPtr<Var> variable;
  const Tensor& value = context->input(1);
  // Note: every resource-variable-manipulating op assumes copy-on-write
  // semantics, and creates a copy of the variable's Tensor if its refcount is
  // bigger than 1 when we try to modify it. This means we never need to copy
  // the original tensor for AssignVariableOp; even if there are other live

                                    std::move(patterns)))) {
    getOperation().emitError("Composite lowering pass failed.");
    signalPassFailure();
  }
}

}  // namespace

// Creates an instance of the pass.
std::unique_ptr<OperationPass<ModuleOp>> CreateCompositeLoweringPass() {
  return std::make_unique<CompositeLoweringPass>();
}

// Registers the pass implementation