return UnrankedTensorType::get(element_type);
}

// Returns true if the op creates a TensorList.
bool IsTensorListInitOp(Operation* op) {
  return isa<TensorListReserveOp>(op) || isa<EmptyTensorListOp>(op) ||
         isa<TensorListFromTensorOp>(op);
}

// Returns the `element_shape` operand of the ops that create a TensorList.
Value GetElementShapeOperand(Operation* op) {

  return strings::StrCat(forward_frame_name, "_backprop");
}

// Creates a loop that counts the number of iterations performed by the
// while loop associated with `while_ctx`. The returned output yields the
// iteration count.
Status AddForwardLoopCounter(WhileContext* while_ctx, const Scope& scope,
                             Output* count) {
  // Create while loop:
  //   i = 0
  //   while forward loop predicate is true:

      if (elem_type && elem_type.hasStaticShape()) return elem_type;
    }
  }
  return std::nullopt;
}

// Creates a ReadVariableOp on a local variable.
Value ReadLocalVariable(Value local_var, OpBuilder builder, Location loc) {
  return builder
      .create<TF::ReadVariableOp>(
          loc,
          ArrayRef<Type>{mlir::cast<TF::ResourceType>(

      // Add dequantize.
      builder.setInsertionPointAfter(read_variable_op);
      auto new_read_variable_op =
          builder.create<ReadVariableOp>(read_variable_op.getLoc(), ref_qtype,
                                         read_variable_op.getResourceId());
      auto new_dq_op = builder.create<DequantizeOp>(
          read_variable_op.getLoc(), read_variable_op.getResult().getType(),
          new_read_variable_op.getResult());

  // available during compilation or compilation only device for on demand
  // execution which may create a recursion if used for constant folding.
  std::string host_cpu = tensorflow::DeviceNameUtils::FullName(
      /*job=*/"localhost", /*replica=*/0, /*task=*/0, /*type=*/"CPU", /*id=*/0);

  absl::StatusOr<OpKernelRunner> runner = OpKernelRunner::Create(
      node_def->get()->op(), node_def->get()->name(), host_cpu, operands.size(),

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

      if (operand.getType() != expected_type) {
        operand = rewriter.create<TF::CastOp>(
            op.getLoc(), expected_type, operand,
            /*Truncate=*/rewriter.getBoolAttr(false));
      }
      casted_operands.push_back(operand);
    }

    auto call = rewriter.create<func::CallOp>(
        op->getLoc(), main_fn.getSymName(), main_fn.getResultTypes(),
        casted_operands);

    b.setInsertionPoint(op);
    if (auto dq = llvm::dyn_cast<DequantizeOp>(op)) {
      auto dcast = b.create<quantfork::DequantizeCastOp>(
          dq.getLoc(), dq.getOutput().getType(), dq.getInput());
      dq.getOutput().replaceAllUsesWith(dcast);
      dq.erase();
    } else if (auto q = llvm::dyn_cast<QuantizeOp>(op)) {
      auto qcast = b.create<quantfork::QuantizeCastOp>(
          q.getLoc(), q.getOutput().getType(), q.getInput());