[XLA:CLIENT] Why force using U32 indices type when converting torch gather op

[Nullkooland]

This commit: a1c04b8 intrudoces the code that forces xla::TorchGather to convert indices tensor with element type of i64 to u32:

xla/xla/client/lib/slicing.cc

Lines 148 to 152 in a1c04b8

	if (ShapeUtil::ElementHasBitWidth(index_shape, 64) &&
	input_shape.dimensions(dim) < std::numeric_limits<uint32>::max()) {
	index = ConvertElementType(index, U32);
	index_shape.set_element_type(U32);
	}

This causes the StableHLO IRs exported with torch_xla to have such pattern around stablehlo.gather:

However, when lowering stablehlo.gather to MLIR tensor dialect's tensor.gather, the tensor.gather op requires that its indices operand tensor has signless integer type, so ui32 would cause error. Also, it is inconsistent that any other index type in the IR is i64 while only index type of this gather is ui32.

@blakehechtman Could you have a look? This commit looks like a HACK to me.
Is it possible to revert this?

[cheshire]

The commit is 5 years old, and thus unlikely to be reverted.

However, when lowering stablehlo.gather to MLIR tensor dialect's tensor.gather

This seems outside of the scope of OpenXLA, more like a stablehlo/tensor dialect interop issue?

[Nullkooland]

The commit is 5 years old, and thus unlikely to be reverted.

However, when lowering stablehlo.gather to MLIR tensor dialect's tensor.gather

This seems outside of the scope of OpenXLA, more like a stablehlo/tensor dialect interop issue?

I don't think this is a stablehlo-to-tensor dialect conversion issue, since stablehlo.gather is able to take i64 indices. After removing this force u32 indices conversion code in XLA, the same torch gather op will be exported as:

I don't see why do we need this u32 HACK.