Fixes #160513

The Problem Summary

The issue boiled down to data type promotion logic. The code base has two different functions that deal with dtype promotion logic. If it is purely multi-dimensional tensor operations, the cpp code gets triggered and that follows the numpy dtype promotion logic. That is why in #160513 NDim tensors are fine as NDim dtypes gets precedence. The issue came with python scalars and 0Dim tensors. When it detects "scalars", a python implementation of dtype promotion logic gets triggered (torch/_prims_common/init.py:1544). Since this is in python, the implementation can't distinguish what is from a wrapped tensor and a 0Dim tensor and thus will just take the highest dtype which is the python double wrapped number.

The Fix

The python implementation for dtype promotion had to know where the scalar came from. Once the scalar can be distinguished then the appropriate dtype can be set. The first approach was to try and expose the is_wrapped_number method but this came with a big issue. During the forward_ad the derivative of those scalars turned out to be ZeroTensors. The ZeroTensor internally uses a hack to initialize a meta dtype tensor which skips expensive dispatch operations. But the copy would not grab everything especially the is_number_wrapped_ property. I thought about modifying the copy but that seemed to go away from the spirit of what the copy was intended for and plus the tests for is_wrapped_number_ requires dim > 0 and a scalar ZeroTensor is a meta dtype tensor which complicates things.

So I chose the route of creating a new property called was_wrapped_number and exposed this property to the python tensor API. I had to modify the autograd code generation to set was_wrapped_number in the mul, add, and div operations in VariableType.cpp. Once this property was set, the dtype promotion logic could be updated to consider wrapped numbers and 0Dim numbers. Once that hierarchy was taken care of, the buggy behavior was fixed.

I wrote a new ops testing module TestForwardADWithScalars. I saw that this bug was unique and required new testing paradigm. This only tests the multiply, add, and divide and I chose this because all operations boil down to these three operations.

[edit]: Just used efficientzerotensor meta and converted that to a python number. Since wrapped number is converted back to a python number, dtype promotion is preserved. The constraint to achieve this happened by setting the forward grad zero tensor of a wrapped number with a wrapped number flag since the tangent of the wrapped number should still be a wrapped number. After that this specific zerotensor was then sent through as a meta type in the BinaryOps.cpp to get appropriate dtype for resulting arithmetic.

@ezyang @OihanJoyot

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel