[JackCaoG]

This will fix #4155

TODO: verify if getDimensionSize actually return a s64, if not I might need to do a cast to s64.

FYI @Krovatkin @vanbasten23

[JackCaoG]

Hmm, ok. This fix seems incorrect because setDimensionSize and getDimensionSize take/returns s32 for TPU. This means TPU expects index and size to be s32

(Pdb) torch_xla._XLAC._get_xla_tensor_dimension_size(t2, 0)
tensor(10, device='xla:0', dtype=torch.int32)

this conflicts with the pytorch assumption where index and size should be s64

(Pdb) torch.nonzero(t1.cpu()).dtype
torch.int64

While I can cast the s32 to s64 on TPU, it doesn't make much sense TPU just can't represent s64 size and I doubt if we ever encounter a dimension that's larger than s32-max. My next step would be try to force non_zero on TPU also returns a logical element type of s32

[JackCaoG]

confirmed that in HLO, get-dimension-size returns a s32

(Pdb) size = torch_xla._XLAC._get_xla_tensor_dimension_size(t2, 0)
(Pdb) print(torch_xla._XLAC._get_xla_tensors_hlo([size]))
HloModule IrToHlo.4, entry_computation_layout={(s64[<=10,2]{0,1})->(s32[])}

ENTRY %IrToHlo.4 (p0.1: s64[<=10,2]) -> (s32[]) {
  %p0.1 = s64[<=10,2]{0,1} parameter(0)
  %get-dimension-size.2 = s32[] get-dimension-size(s64[<=10,2]{0,1} %p0.1), dimensions={0}
  ROOT %tuple.3 = (s32[]) tuple(s32[] %get-dimension-size.2)
}

[JackCaoG]

ok now on TPU, nonzero's real xla type and dtype will be s32, I am sure some test will fail on TPU.. Let me give it a try.

[vanbasten23]

Quick question: in your investigation, why did you do (Pdb) torch_xla._XLAC._get_xla_tensor_dimension_size(t2, 0) as opposed to t2.dtype? Is it because torch_xla._XLAC._get_xla_tensor_dimension_size returns the size from XLA layer?

[JackCaoG]

dtype was overwritten by logical_element_type

xla/torch_xla/csrc/tensor.cpp

Lines 552 to 556 in f858155

	at::ScalarType XLATensor::dtype() const {
	return data()->logical_element_type
	? *data()->logical_element_type
	: TensorTypeFromXlaType(shape().get().element_type());
	}

[vanbasten23]

mind adding a comment about where the cast happens, such as #4155 (comment).?

[vanbasten23]

I assume what this CreateFrom does is: input->GetIrValue()(line1878) contains the IR graph before we add the NonZero node; then input->CreateFrom insert the NonZero node to the existing IR graph. Is my understanding correct?

[JackCaoG]

CreateFrom will create a new XLATensor using the IR node provided but it will use the same logical_element_type and device from the input tensor. I used Create here instead.

[JackCaoG]

@vanbasten23 Can you test if dynamic input liner model will works on PJRT + TPU? @will-cromar had a bug open for supporting dynamic shape + PJRT so we want to figure out the gap here.

[vanbasten23]

@vanbasten23 Can you test if dynamic input liner model will works on PJRT + TPU? @will-cromar had a bug open for supporting dynamic shape + PJRT so we want to figure out the gap here.

How can I test it on PJRT+TPU? Is it by setting PJRT_DEVICE=TPU and run the model on my TPU VM?

[will-cromar]

Right. All you need to do is set PJRT_DEVICE=TPU on TPU VM https://github.com/pytorch/xla/blob/master/docs/pjrt.md

[JackCaoG]

I am able to repo the GPU failure with

XLA_EXPERIMENTAL="nonzero:masked_select" ./build/test_ptxla --gtest_filter=AtenXlaTensorTest.TestNonzero

and the error is expected

Different shape:
long [3, 2]
-vs-
int [3, 2]

I guess I will cast both to s64 before doing the allClose

[JackCaoG]

adding a _cast_Long seems to resolve the error

AllClose(b, torch::_cast_Long(xla_b));

I will post a fix.

[vanbasten23]

@vanbasten23 Can you test if dynamic input liner model will works on PJRT + TPU? @will-cromar had a bug open for supporting dynamic shape + PJRT so we want to figure out the gap here.

How can I test it on PJRT+TPU? Is it by setting PJRT_DEVICE=TPU and run the model on my TPU VM?

I tested the forward pass model on PJRT+TPU and it works fine.

The model returns a different loss each time I run. But this problem exists for both XRT and PJRT, so I think it might be an issue with the model itself. @JackCaoG @will-cromar

[JackCaoG]

hmm @vanbasten23 if we fixed the input we should see the same result every time. Can you look into why that's not the case? We need to confirm that dynamic_shape on xla return the same result as on CPU.

[vanbasten23]

hmm @vanbasten23 if we fixed the input we should see the same result every time. Can you look into why that's not the case? We need to confirm that dynamic_shape on xla return the same result as on CPU.

Yeah, I'm already looking into it. Will let you know if I find anything.

[vanbasten23]

Not directly related to this pr but do you know why NonZero has 2 outputs:

xla/torch_xla/csrc/ops/nonzero.cpp

Line 27 in d668f2c

/*num_outputs=*/2) {}

?

[JackCaoG]

Not directly related to this pr but do you know why NonZero has 2 outputs:

xla/torch_xla/csrc/ops/nonzero.cpp

Line 27 in d668f2c

/*num_outputs=*/2) {}

?

if you look at the lowering, it returns 2 XLAOP. However second XLAOP was discarded. I am not sure why that's the case.

[vanbasten23]

hmm @vanbasten23 if we fixed the input we should see the same result every time. Can you look into why that's not the case? We need to confirm that dynamic_shape on xla return the same result as on CPU.

Yeah, I'm already looking into it. Will let you know if I find anything.

I figured out why the loss varies each time. By default, the weights in each layer has some randomization. So to make the loss deterministic and make it easy to write test, I need to initialize the weights explicitly by myself. I've updated the model here.

[vanbasten23]

@JackCaoG in the first comment of this pr, you mentioned TODO: verify if getDimensionSize actually return a s64, if not I might need to do a cast to s64. I wonder if this TODO is still valid.

I ran the test in my pr and dump the HLO graph:

%set-dimension-size.39 = s32[<=10,2]{1,0} set-dimension-size(s32[10,2]{1,0} %concatenate.38, s32[] %reduce.17), dimensions={0}
..
%get-dimension-size.130 = s32[] get-dimension-size(f32[<=10]{0} %multiply.123), dimensions={0}

It looks getDimensionSize returns s32.

[JackCaoG]

yea, I decided to make the SizeNode to be s32 on TPU instead, so we are honest about the real underlying type and avoid additional casting.