for (auto assign_variable_op : llvm::make_early_inc_range(
           session_init_func.getOps<TF::AssignVariableOp>())) {
    Value resource_operand = assign_variable_op.getOperand(0);
    Value assigned_value_operand = assign_variable_op.getOperand(1);

    if (auto var_handle_op =
            dyn_cast<TF::VarHandleOp>(resource_operand.getDefiningOp());
        var_handle_op &&

           session_init_func.getOps<TF::AssignVariableOp>())) {
    Value resource_operand = assign_variable_op.getOperand(0);
    auto var_handle_op =
        dyn_cast<TF::VarHandleOp>(resource_operand.getDefiningOp());
    if (!var_handle_op) continue;

    Value assigned_value_operand = assign_variable_op.getOperand(1);
    auto const_op =
        dyn_cast<TF::ConstOp>(assigned_value_operand.getDefiningOp());

 public:
  using OpRewritePattern<mlir::stablehlo::TransposeOp>::OpRewritePattern;

  LogicalResult match(mlir::stablehlo::TransposeOp op) const override {
    Value operand = op.getOperand();
    auto const_op =
        dyn_cast_or_null<mlir::stablehlo::ConstantOp>(operand.getDefiningOp());
    if (!const_op) return failure();

    // Only support float tensors.

  ShapedType type;
  if (llvm::isa<TF::ReadVariableOp>(op)) {
    type = op->getResult(0).getType().cast<ShapedType>();
  } else if (llvm::isa<TF::AssignVariableOp>(op)) {
    type = op->getOperand(1).getType().cast<ShapedType>();
  } else if (llvm::isa<TF::VarHandleOp>(op)) {
    type = op->getResult(0)
               .getType()
               .cast<ShapedType>()
               .getElementType()

  output_shape[1] = composite_result_shape[2];
  output_shape[2] = composite_result_shape[3];
  output_shape[3] = composite_result_shape[1];

  auto input_type = mlir::cast<ShapedType>(old_op->getOperand(0).getType());

  return RankedTensorType::get(output_shape, input_type.getElementType());
}
}  // namespace odml

  // result all correspond to each other (see definition of `WhileRegionOp`).
  int idx = result.getResultNumber();
  Value body_yield_operand = body_yield_op->getOperand(idx);
  Value body_block_argument = body_block.getArgument(idx);
  Value cond_block_argument = cond_block.getArgument(idx);
  // Consider the op that defines the unused result as a candidate for pruning.

}

// `is_cpu_read` is set to `true` iff `read` is on a resource with device type
// CPU.
LogicalResult IsCpuRead(FuncOp func, ReadVariableOp read, bool& is_cpu_read) {
  if (auto arg = mlir::dyn_cast<BlockArgument>(read->getOperand(0))) {
    if (arg.getOwner() != &(func.front())) {
      is_cpu_read = false;
      return success();
    }
    if (auto attr = func.getArgAttrOfType<StringAttr>(arg.getArgNumber(),

  std::vector<int> stateful_operands_index;
  if (!IsStatefulOp(op, &stateful_operands_index)) return success();

  for (int index : stateful_operands_index) {
    Value operand = op->getOperand(index);
    auto inserted_value = values->insert(operand).second;
    if (inserted_value) continue;
    // We can only clone the constant op or const->dequantize combo. The latter

      // parallel_execute region.
      assign_var.getOperation()->moveBefore(terminator);
      assign_var.getValueMutable().assign(
          terminator->getOperand(result.index()));
      results_to_remove.push_back(result.index());
    }

    rewrite |= !results_to_remove.empty();
  }

  if (!rewrite) return;

  LogicalResult matchAndRewrite(mlir::stablehlo::UniformQuantizeOp op,
                                PatternRewriter& rewriter) const override {
    DenseFPElementsAttr attr;
    if (matchPattern(op.getOperand(), m_Constant(&attr))) {
      const Type qtype = op.getResult().getType();
      ElementsAttr quantized_attr = Quantize(attr, qtype);
      if (quantized_attr) {