Align `_choose_qparams_affine ` with `_choose_scale_float8`

[jerryzh168]

In _choose_qparams_affine (for int):

we use keepdim=False:

ao/torchao/quantization/quant_primitives.py

Lines 1552 to 1553 in b4ec4cb

	min_val = torch.amin(input, dim=reduction_dims, keepdim=False)
	max_val = torch.amax(input, dim=reduction_dims, keepdim=False)

that will make per tensor quantization to have a scalar scale automatically.

I think we can change _choose_qparams_affine to use keepdim=True and make sure we match the scale/zero_point with dimensions of input, so we don't need to do this:

ao/torchao/quantization/quantize_/workflows/intx/intx_unpacked_to_int8_tensor.py

Lines 247 to 254 in 0ffbac1

	# Reshape scale and zero_point to be compatible with block_size
	# This is asserted in IntxUnpackedToInt8Tensor's __init__
	n_blocks = []
	for i in range(len(block_size)):
	assert qdata.shape[i] % block_size[i] == 0
	n_blocks.append(qdata.shape[i] // block_size[i])
	scale = scale.reshape(*n_blocks)
	zero_point = zero_point.reshape(*n_blocks)

we could also remove the block_size argument from quantize_affine and dequantize_affine afterwards since the rank are already aligned.

we can also add some docs afterwards to both ops to clarify.

[jerryzh168]

@namgyu-youn I updated this issue a bit, I think we can actually align the integer choose_qparams_affine with the float8 one, and keep the original rank instead of shrink it, this would make the scale/zero_point easier to handle overall

created a separate issue for float8: #3414

[AryanBagade]

@jerryzh168 if this issue is open for external contributors, shall i work on it?

[jerryzh168]

yes, please go ahead, thanks for helping! we also have internal tests that maybe affected so I can help with testing that when you put up the PR

[AryanBagade]

@jerryzh168 created a PR, thanks for letting me contribute, once this is merged will work on its twin issue !

[namgyu-youn]

@namgyu-youn I updated this issue a bit, I think we can actually align the integer choose_qparams_affine with the float8 one, and keep the original rank instead of shrink it, this would make the scale/zero_point easier to handle overall

created a separate issue for float8: #3414

Thanks so much for the info Jerry. I was too busy to work on this, and hope it will be resolved by @AryanBagade .

btw, there are a few PRs not merged after approvals (#3143, #3148, #3331). Can you please look into it? I really don't want to ping for distraction, but delayed PRs are making it hard for me to work more.