weight_only_ptq.input_quantized_types().at(1);
    if (quantized_type.dimension_specs().has_dimension()) {
      return quantized_type.dimension_specs().dimension();
    }
    return GetDefaultQuantizationDimension(op);
  }

  // Determines quantization dimension of weights for given `tf.XlaCallModule`
  // op. For convolution, returns output feature dimension of the kernel. For