op_operand.set(builder.create<TF::CastOp>(
            op->getLoc(), ToBfloat16Type(op_operand.get().getType()),
            op_operand.get()));
      }
    }
    builder.setInsertionPointAfter(op);
    for (auto op_result : op->getOpResults()) {
      if (IsLargeFloatType(op_result.getType())) {
        const Type original_type = op_result.getType();
        op_result.setType(ToBfloat16Type(original_type));

      quant_dim = input_type.getRank() - 1;
    }
    // Use the min/max from the operands and the num_bits and narrow_range
    // attribute to create the quantization parameter for the new quantize op.
    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();

  return true;
}

void QuantizationDriver::QuantizeOpResult(Operation* op, const int result_index,
                                          const QuantizedType quantized_type) {
  builder_.setInsertionPointAfter(op);
  const Value original_result = op->getResult(result_index);
  QuantizeValue(original_result, quantized_type, op->getLoc());
}

void QuantizationDriver::QuantizeArg(BlockArgument arg,

        return failure();
      }
    }
    if (call_op->use_empty()) return success();

    DenseMap<Value, quantfork::StorageCastOp> replace_map;
    rewriter.setInsertionPointAfter(call_op);

    SmallVector<Type, 4> result_types;
    for (Value result : call_op->getResults()) {
      TensorType result_type = mlir::dyn_cast<TensorType>(result.getType());
      if (!result_type) {

    }
    // Use the min/max from the operands and the num_bits and narrow_range
    // attribute to create the quantization parameter for the new quantize op.
    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();

    }
    // Use the min/max from the operands and the num_bits and narrow_range
    // attribute to create the quantization parameter for the new quantize op.
    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();

                             Type new_user_op_type) const {
    auto op_before_dequantize = original_dequantize_op.getOperand(0);

    // Create a new dequantize op that is propagated.
    rewriter.setInsertionPointAfter(user_op);
    TF::PartitionedCallOp new_dequantize_op =
        cast<TF::PartitionedCallOp>(rewriter.clone(*original_dequantize_op));

    // Skip the original dequant op and connect the op before dequantize to the

      auto read_variable_op = dyn_cast_or_null<ReadVariableOp>(var_handle_user);
      if (!read_variable_op) continue;
      // 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());

        gemm_style_shape, accumulation_quantized_element_type);
  }

  gemm_style_op_result.setType(new_gemm_style_op_result_type);

  rewriter.setInsertionPointAfter(gemm_style_op);

  Operation* next_op = FindUserOfType<>(gemm_style_op);

  // If activation exists, omit clipping op.
  // Since out_scale and out_zp are computed based on clipped range,

      if (legacy_float_scale) {
        quant_type = DownCastScale(quant_type, rmin, rmax, op->getLoc());
      }
    } else {
      return failure();
    }

    rewriter.setInsertionPointAfter(op.getOperation());
    Type result_type = quant_type.castFromExpressedType(op.getType());
    auto q =
        rewriter.create<QuantizeOpT>(op.getLoc(), result_type, op.getArg());