Dynamic `sym_sizes_custom()` fails at `.test_reference_eager_median_xla_int64`

[miladm]

🐛 Bug

The dynamic version of sym_sizes_custom causes size mismatch failures.

======================================================================
FAIL: test_reference_eager_median_xla_int64 (__main__.TestOpInfoXLA)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/torch/testing/_internal/common_device_type.py", line 377, in instantiated_test
    result = test(self, **param_kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/testing/_internal/common_device_type.py", line 814, in test_wrapper
    return test(*args, **kwargs)
  File "/workspace/pytorch/xla/test/test_ops.py", line 411, in test_reference_eager
    self.compare_with_eager_reference(op, sample_input)
  File "/workspace/pytorch/xla/test/test_ops.py", line 402, in compare_with_eager_reference
    self.assertEqual(actual, expected, exact_dtype=True, exact_device=False)
  File "/workspace/pytorch/xla/test/pytorch_test_base.py", line 629, in assertEqual
    return DeviceTypeTestBase.assertEqual(self, x, y, *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/torch/testing/_internal/common_utils.py", line 2392, in assertEqual
    msg=(lambda generated_msg: f"{generated_msg} : {msg}") if isinstance(msg, str) and self.longMessage else msg,
  File "/opt/conda/lib/python3.7/site-packages/torch/testing/_comparison.py", line 1093, in assert_equal
    raise error_metas[0].to_error(msg)
AssertionError: The values for attribute 'shape' do not match: torch.Size([0]) != torch.Size([]).

The failure occurred for item [0]

[miladm]

Looking into the tests, it seems the error is due to the following two torch_fn calls:

torch.return_types.median(values=tensor(-6, device='xla:0'),indices=tensor(0, device='xla:0'))
torch.return_types.median(values=tensor(-6),indices=tensor(0))

[miladm]

More insights into this failure and how it occurs:

The test_ops script, produces the following two sample_inputs: SampleInput(input=-6, args=(), kwargs={'dim': 0, 'keepdim': True} and SampleInput(input=-6, args=(), kwargs={}

When dim=0 and keepdim=True are set, the test fails. Otherwise, the test passes. In both test cases, the input tensor is torch.tensor(-6, device='xle:0') and torch.tensor(-6). However, the output size from torch_xla is torch.Size([0]) incorrectly.

Looking at the C++ PR implementation, I confirmed the source of the failure is tensor_impl constructor detects a rank 1 for a tensor that's otherwise expected to have rank 0. As a result, the test fails. I looked at the IR dump to determine what op created the incorrectly shaped tensor; though, the IR dump hasn't helped isolate the op yet. The dump is included below for reference.

IR {
  %0 = s64[] xla::device_data(), device=CPU:0
  %1 = s64[] prim::Constant(), value=214013
  %2 = s64[] aten::mul(%1, %0)
  %3 = s64[] prim::Constant(), value=2531011
  %4 = s64[] aten::add(%3, %2)
  %5 = s64[] xla::device_data(), device=CPU:0
  %6 = s64[] xla::device_data(), device=CPU:0
  %7 = s64[] aten::random(%6, %5, %4), ROOT=0
}

IR {
  %0 = s64[] xla::device_data(), device=CPU:0
  %1 = s64[] prim::Constant(), value=214013
  %2 = s64[] aten::mul(%1, %0)
  %3 = s64[] prim::Constant(), value=2531011
  %4 = s64[] aten::add(%3, %2)
  %5 = s64[] xla::device_data(), device=CPU:0
  %6 = s64[] xla::device_data(), device=CPU:0
  %7 = s64[] aten::random(%6, %5, %4), ROOT=0
}

IR {
  %0 = s64[] xla::device_data(), device=CPU:0
  %1 = s64[] prim::Constant(), value=214013
  %2 = s64[] aten::mul(%1, %0)
  %3 = s64[] prim::Constant(), value=2531011
  %4 = s64[] aten::add(%3, %2)
  %5 = s64[] xla::device_data(), device=CPU:0
  %6 = s64[] xla::device_data(), device=CPU:0
  %7 = s64[] aten::random(%6, %5, %4), ROOT=0
}

IR {
  %0 = s64[] prim::Constant(), value=0
  %1 = s64[0]{0} aten::expand(%0), size=(0), ROOT=0
  %2 = s64[] prim::Constant(), value=0
  %3 = s64[0]{0} aten::expand(%2), size=(0), ROOT=1
}

Wild theory: I am looking to learn if the rank returns 1 because the stack must have checked and sliced over a dim field before hitting tensor_impl.

[miladm]

Below is the HLO dump for reference:

HloModule IrToHlo.28

ENTRY %IrToHlo.28 (p0.1: s64[], p1.6: s64[], p2.7: s64[]) -> (s64[]) {
  %constant.4 = s64[] constant(2531011)
  %constant.2 = s64[] constant(214013)
  %p0.1 = s64[] parameter(0)
  %multiply.3 = s64[] multiply(s64[] %constant.2, s64[] %p0.1)
  %add.5 = s64[] add(s64[] %constant.4, s64[] %multiply.3)
  %convert.8 = u64[] convert(s64[] %add.5)
  %reshape.10 = u64[1]{0} reshape(u64[] %convert.8)
  %constant.9 = u64[] constant(0)
  %reshape.11 = u64[1]{0} reshape(u64[] %constant.9)
  %concatenate.12 = u64[2]{0} concatenate(u64[1]{0} %reshape.10, u64[1]{0} %reshape.11), dimensions={0}
  %rng-bit-generator.13 = (u64[2]{0}, u64[]) rng-bit-generator(u64[2]{0} %concatenate.12), algorithm=rng_default
  %get-tuple-element.15 = u64[2]{0} get-tuple-element((u64[2]{0}, u64[]) %rng-bit-generator.13), index=0
  %p2.7 = s64[] parameter(2)
  %get-tuple-element.14 = u64[] get-tuple-element((u64[2]{0}, u64[]) %rng-bit-generator.13), index=1
  %p1.6 = s64[] parameter(1)
  %bitcast-convert.17 = u64[] bitcast-convert(s64[] %p1.6)
  %bitcast-convert.16 = u64[] bitcast-convert(s64[] %p2.7)
  %subtract.18 = u64[] subtract(u64[] %bitcast-convert.17, u64[] %bitcast-convert.16)
  %remainder.19 = u64[] remainder(u64[] %get-tuple-element.14, u64[] %subtract.18)
  %constant.20 = u64[] constant(1)
  %shift-right-logical.21 = u64[] shift-right-logical(u64[] %remainder.19, u64[] %constant.20)
  %bitcast-convert.24 = s64[] bitcast-convert(u64[] %shift-right-logical.21)
  %add.25 = s64[] add(s64[] %p2.7, s64[] %bitcast-convert.24)
  %subtract.22 = u64[] subtract(u64[] %remainder.19, u64[] %shift-right-logical.21)
  %bitcast-convert.23 = s64[] bitcast-convert(u64[] %subtract.22)
  %add.26 = s64[] add(s64[] %add.25, s64[] %bitcast-convert.23)
  ROOT %tuple.27 = (s64[]) tuple(s64[] %add.26)
}

ENTRY %IrToHlo.28 (p0.1: s64[], p1.6: s64[], p2.7: s64[]) -> (s64[]) {
  %constant.4 = s64[] constant(2531011)
  %constant.2 = s64[] constant(214013)
  %p0.1 = s64[] parameter(0)
  %multiply.3 = s64[] multiply(s64[] %constant.2, s64[] %p0.1)
  %add.5 = s64[] add(s64[] %constant.4, s64[] %multiply.3)
  %convert.8 = u64[] convert(s64[] %add.5)
  %reshape.10 = u64[1]{0} reshape(u64[] %convert.8)
  %constant.9 = u64[] constant(0)
  %reshape.11 = u64[1]{0} reshape(u64[] %constant.9)
  %concatenate.12 = u64[2]{0} concatenate(u64[1]{0} %reshape.10, u64[1]{0} %reshape.11), dimensions={0}
  %rng-bit-generator.13 = (u64[2]{0}, u64[]) rng-bit-generator(u64[2]{0} %concatenate.12), algorithm=rng_default
  %get-tuple-element.15 = u64[2]{0} get-tuple-element((u64[2]{0}, u64[]) %rng-bit-generator.13), index=0
  %p2.7 = s64[] parameter(2)
  %get-tuple-element.14 = u64[] get-tuple-element((u64[2]{0}, u64[]) %rng-bit-generator.13), index=1
  %p1.6 = s64[] parameter(1)
  %bitcast-convert.17 = u64[] bitcast-convert(s64[] %p1.6)
  %bitcast-convert.16 = u64[] bitcast-convert(s64[] %p2.7)
  %subtract.18 = u64[] subtract(u64[] %bitcast-convert.17, u64[] %bitcast-convert.16)
  %remainder.19 = u64[] remainder(u64[] %get-tuple-element.14, u64[] %subtract.18)
  %constant.20 = u64[] constant(1)
  %shift-right-logical.21 = u64[] shift-right-logical(u64[] %remainder.19, u64[] %constant.20)
  %bitcast-convert.24 = s64[] bitcast-convert(u64[] %shift-right-logical.21)
  %add.25 = s64[] add(s64[] %p2.7, s64[] %bitcast-convert.24)
  %subtract.22 = u64[] subtract(u64[] %remainder.19, u64[] %shift-right-logical.21)
  %bitcast-convert.23 = s64[] bitcast-convert(u64[] %subtract.22)
  %add.26 = s64[] add(s64[] %add.25, s64[] %bitcast-convert.23)
  ROOT %tuple.27 = (s64[]) tuple(s64[] %add.26)
}

ENTRY %IrToHlo.28 (p0.1: s64[], p1.6: s64[], p2.7: s64[]) -> (s64[]) {
  %constant.4 = s64[] constant(2531011)
  %constant.2 = s64[] constant(214013)
  %p0.1 = s64[] parameter(0)
  %multiply.3 = s64[] multiply(s64[] %constant.2, s64[] %p0.1)
  %add.5 = s64[] add(s64[] %constant.4, s64[] %multiply.3)
  %convert.8 = u64[] convert(s64[] %add.5)
  %reshape.10 = u64[1]{0} reshape(u64[] %convert.8)
  %constant.9 = u64[] constant(0)
  %reshape.11 = u64[1]{0} reshape(u64[] %constant.9)
  %concatenate.12 = u64[2]{0} concatenate(u64[1]{0} %reshape.10, u64[1]{0} %reshape.11), dimensions={0}
  %rng-bit-generator.13 = (u64[2]{0}, u64[]) rng-bit-generator(u64[2]{0} %concatenate.12), algorithm=rng_default
  %get-tuple-element.15 = u64[2]{0} get-tuple-element((u64[2]{0}, u64[]) %rng-bit-generator.13), index=0
  %p2.7 = s64[] parameter(2)
  %get-tuple-element.14 = u64[] get-tuple-element((u64[2]{0}, u64[]) %rng-bit-generator.13), index=1
  %p1.6 = s64[] parameter(1)
  %bitcast-convert.17 = u64[] bitcast-convert(s64[] %p1.6)
  %bitcast-convert.16 = u64[] bitcast-convert(s64[] %p2.7)
  %subtract.18 = u64[] subtract(u64[] %bitcast-convert.17, u64[] %bitcast-convert.16)
  %remainder.19 = u64[] remainder(u64[] %get-tuple-element.14, u64[] %subtract.18)
  %constant.20 = u64[] constant(1)
  %shift-right-logical.21 = u64[] shift-right-logical(u64[] %remainder.19, u64[] %constant.20)
  %bitcast-convert.24 = s64[] bitcast-convert(u64[] %shift-right-logical.21)
  %add.25 = s64[] add(s64[] %p2.7, s64[] %bitcast-convert.24)
  %subtract.22 = u64[] subtract(u64[] %remainder.19, u64[] %shift-right-logical.21)
  %bitcast-convert.23 = s64[] bitcast-convert(u64[] %subtract.22)
  %add.26 = s64[] add(s64[] %add.25, s64[] %bitcast-convert.23)
  ROOT %tuple.27 = (s64[]) tuple(s64[] %add.26)
}

ENTRY %IrToHlo.12 () -> (s64[0], s64[0]) {
  %constant.1 = s64[] constant(0)
  %reshape.2 = s64[1]{0} reshape(s64[] %constant.1)
  %broadcast.3 = s64[1]{0} broadcast(s64[1]{0} %reshape.2), dimensions={0}
  %reshape.4 = s64[] reshape(s64[1]{0} %broadcast.3)
  %broadcast.5 = s64[0]{0} broadcast(s64[] %reshape.4), dimensions={}
  %constant.6 = s64[] constant(0)
  %reshape.7 = s64[1]{0} reshape(s64[] %constant.6)
  %broadcast.8 = s64[1]{0} broadcast(s64[1]{0} %reshape.7), dimensions={0}
  %reshape.9 = s64[] reshape(s64[1]{0} %broadcast.8)
  %broadcast.10 = s64[0]{0} broadcast(s64[] %reshape.9), dimensions={}
  ROOT %tuple.11 = (s64[0]{0}, s64[0]{0}) tuple(s64[0]{0} %broadcast.5, s64[0]{0} %broadcast.10)
}