Repository: microsoft/onnxruntime
File: onnxruntime/core/optimizer/qdq_transformer/weight_bias_quantization.cc (C++)

Code before:
    // Nothing to do if neither weight nor bias is initializer.
    if (dq_1 && bias_size == -1) {
      continue;
    }

    bool is_per_tensor_scale = true;
    // If weight is quantized, it's either per-tensor or per-channel on specific axis (0 for Conv and 1 for Gemm).
    if (dq_1 && !optimizer_utils::IsScalar(*dq_1->InputDefs()[1])) {
      is_per_tensor_scale = false;
      auto weight_scale_shape = dq_1->InputDefs()[1]->Shape();
      if (!weight_scale_shape || weight_scale_shape->dim_size() != 1 || !weight_scale_shape->dim(0).has_dim_value() ||
          weight_scale_shape->dim(0).dim_value() != bias_size) {
        continue;
      }

      const auto& dq_attrs = dq_1->GetAttributes();
      auto attr_it = dq_attrs.find("block_size");
      // Default value of block_size=0 has no significance. Don't skip weight_bias_quantization.
      if (attr_it != dq_attrs.end() && attr_it->second.i() != 0) {
        continue;
      }

      int64_t axis = 1;
      if (auto axis_iter = dq_attrs.find("axis"); axis_iter != dq_attrs.end()) {
        const ONNX_NAMESPACE::TensorShapeProto* weight_shape = weight_arg->Shape();
        int64_t weight_rank = narrow<int64_t>(weight_shape->dim_size());
        axis = HandleNegativeAxis(axis_iter->second.i(), weight_rank);
      }

      int64_t expected_axis = 0;
      if (node.OpType() == "Gemm") {
        int64_t transB = 0;
        const auto& gemm_attrs = node.GetAttributes();
        if (auto trans_b_iter = gemm_attrs.find("transB"); trans_b_iter != gemm_attrs.end()) {
          transB = trans_b_iter->second.i();
        }
        expected_axis = transB == 0 ? 1 : 0;
      }

      if (axis != expected_axis) {
        continue;
      }
    }

    NodeArg* weight_scale_arg = nullptr;
    if (!dq_1) {
      Initializer initializer(graph, *weight_proto, graph.ModelPath());
      const float* weight_data = initializer.data<float>();

      // Quantize float32 weight to int8_t (per-tensor, symmetric).
      // int8_t quantization of input[1] works with input[0] of all types.

Review comment (bug):
I believe protobuf `x_size()` returns `int`, right? that doesn't need a call to `narrow` to convert it to `int64_t`.
```suggestion
        int64_t weight_rank = weight_shape->dim_size();
```

also, is it possible for `weight_shape` to be nullptr?

Code after:
    // Nothing to do if neither weight nor bias is initializer.
    if (dq_1 && bias_size == -1) {
      continue;
    }

    bool is_per_tensor_scale = true;
    // If weight is quantized, it's either per-tensor or per-channel on specific axis (0 for Conv and 1 for Gemm).
    if (dq_1 && !optimizer_utils::IsScalar(*dq_1->InputDefs()[1])) {
      is_per_tensor_scale = false;
      auto weight_scale_shape = dq_1->InputDefs()[1]->Shape();
      if (!weight_scale_shape || weight_scale_shape->dim_size() != 1 || !weight_scale_shape->dim(0).has_dim_value() ||
          weight_scale_shape->dim(0).dim_value() != bias_size) {
        continue;
      }

      const auto& dq_attrs = dq_1->GetAttributes();
      auto attr_it = dq_attrs.find("block_size");
      // Default value of block_size=0 has no significance. Don't skip weight_bias_quantization.
      if (attr_it != dq_attrs.end() && attr_it->second.i() != 0) {
        continue;
      }

      int64_t axis = 1;
      if (auto axis_iter = dq_attrs.find("axis"); axis_iter != dq_attrs.end()) {
        axis = axis_iter->second.i();
        const ONNX_NAMESPACE::TensorShapeProto* weight_shape = weight_arg->Shape();
        if (!weight_shape && dq_1->InputDefs()[0]) {
          weight_shape = dq_1->InputDefs()[0]->Shape();
        }
        if (axis < 0 && !weight_shape) {
          continue;
        }
        axis = HandleNegativeAxis(axis, weight_shape->dim_size());
      }

      int64_t expected_axis = 0;
      if (node.OpType() == "Gemm") {
        int64_t transB = 0;
        const auto& gemm_attrs = node.GetAttributes();
        if (auto trans_b_iter = gemm_attrs.find("transB"); trans_b_iter != gemm_attrs.end()) {
          transB = trans_b_iter->second.i();
        }
        expected_axis = transB == 0 ? 1 : 0;
      }

      if (axis != expected_axis) {
        continue;
      }
    }

    NodeArg* weight_scale_arg = nullptr;