torch.jit.trace appears to convert integer division to tensor division

[StephenHogg]

🐛 Bug

As per this issue in einops, it appears that torch.jit.trace is converting integer division to tensor division.

To Reproduce

Steps to reproduce the behavior:

import torch
import einops

def f(x):
    return x.shape[0] // 4
test_tensor = torch.rand([1, 2, 376, 16, 16])
res = torch.jit.trace(f, test_tensor)

The warning provided is as follows:

/usr/local/lib/python3.7/dist-packages/torch/_tensor.py:575: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.
To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at  /pytorch/aten/src/ATen/native/BinaryOps.cpp:467.)
  return torch.floor_divide(self, other)

Expected behavior

The division in that function shouldn't be getting converted to tensor division.

Environment

PyTorch version: 1.9.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.0 (default, Jun  7 2021, 13:53:36)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.4.72-microsoft-standard-WSL2-x86_64-with-glibc2.29
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.1
[pip3] torch==1.9.0
[conda] Could not collect

cc @gmagogsfm

[arogozhnikov]

Thanks for raising this up, @StephenHogg.

Einops relies on floor division for computation of axes dimensions, and there is no way to change it.

Jit would be broken thus for any pytorch code operating with axis manipulations.

[arogozhnikov]

Hey torch team, it worth implementing proper translation for // unless you want to break jit for ~900 repositories (likely many more).

[gmagogsfm]

Hi @arogozhnikov,

Thanks for reporting this issue.

jit.trace doesn't trace non-tensor (like integer) values. It is recommended to use torch.jit.script to compile any portion of program that has non-tensor values.

Additionally, currently implementation of floor_div in PyTorch is buggy, if you'd like true floor_div behavior, please use following code snipet:

@torch.jit.script
def f(x):
    return torch.div(torch.tensor([x.shape[0]]), 4, rounding_mode='floor')

[arogozhnikov]

It is recommended to use torch.jit.script to compile any portion of program that has non-tensor values.

Ready to do so, right after it becomes any decent, here is a list of what is missing, and see discussion:
arogozhnikov/einops#115 (comment)

Since I don't expect torch.jit.script to have this level of support, there is currently no solution to problem.

jit.trace doesn't trace non-tensor (like integer) values.

if you analyze the code posted and forget about jit, you'll see it's a well-defined piece of code that operates with integer.
It's in jit-trace world this division IS traced, and mapped to torch's (buggy) floordiv. That's what you see in the example.

So, plain simple: function works, doesn't know and never uses pytorch's buggy implementation, uses only built-in div, and it should be translated to correct counterpart in jit.

Specifying floordiv for shape components sounds doable if there is an appropriate abstraction in internal representation. If not - mapping // to torch.div(..., rounding_mode='floor') during jit-tracing is the right approach.

Risk-thinking, I don't see a case when one wants incorrect behavior of floordiv that doesn't match numpy and built-in - that sounds too crazy.
One release with warning and buggy version, one release with warning and correct version, then just correct version without warning. Pulling out an operation breaks more code.

[RuRo]

@gmagogsfm Can you clarify, why is this issue closed?