Looking for PyTorch op data type promotion and result tensor data type rules

Sujoy_Saraswati · January 19, 2022, 1:49am

Is there any document/guideline for PyTorch ops and data type conversions and promotion rules? There is a document at Tensor Attributes — PyTorch 2.1 documentation, however it does not cover some of the op specific cases. For example, add for int32 arguments returns an int32 dtype tensor, but div returns a float32 tensor -

a = torch.ones(4, dtype=torch.int32)
b = torch.ones(4, dtype=torch.int32)
add = a + b
add.dtype
torch.int32
div = a / b
div.dtype
torch.float32
torch.result_type(a, b) ← This doesn’t return the correct result dtype for div
torch.int32

Similarly, for logical comparator ops like eq, the result tensor is boolean but for in-place eq_, the result tensor type is same as the original. Also, eq_ doesn’t seem to do a dtype conversion -

ai.dtype
torch.int32
b.dtype
torch.float32
c = torch.eq(ai, b) ← this works fine
ai.eq_(b) ← this doesn’t work
Traceback (most recent call last):
File “”, line 1, in
RuntimeError: Expected object of scalar type int but got scalar type float for argument ‘other’

It would be great if anyone can share any document/guideline on the rules covering all these cases.

wconstab · May 12, 2022, 9:16pm

Hey Sujoy,

A couple of helpful resources are:

the table that determines actual promotion behavior:

github.com

pytorch/pytorch/blob/master/c10/core/ScalarType.h#L422

      
        
                  false,
                  "promoteTypes with quantized numbers is not handled yet; figure out what the correct rules should be, offending types: ",
                  toString(a),
                  " ",
                  toString(b));
            }
            
            
// this matrix has to be consistent with
            // AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS undefined is used where we
            // are not sure about the correct value for type promotion.
            static constexpr ScalarType _promoteTypesLookup[static_cast<int>(
                ScalarType::NumOptions)][static_cast<int>(ScalarType::NumOptions)] = {
                /*        u1  i1  i2  i4  i8  f2  f4  f8  c2  c4  c8  b1  q1  q2  q3  bf*/
                /* u1 */ {u1, i2, i2, i4, i8, f2, f4, f8, c2, c4, c8, u1, ud, ud, ud, bf},
                /* i1 */ {i2, i1, i2, i4, i8, f2, f4, f8, c2, c4, c8, i1, ud, ud, ud, bf},
                /* i2 */ {i2, i2, i2, i4, i8, f2, f4, f8, c2, c4, c8, i2, ud, ud, ud, bf},
                /* i4 */ {i4, i4, i4, i4, i8, f2, f4, f8, c2, c4, c8, i4, ud, ud, ud, bf},
                /* i8 */ {i8, i8, i8, i8, i8, f2, f4, f8, c2, c4, c8, i8, ud, ud, ud, bf},
                /* f2 */ {f2, f2, f2, f2, f2, f2, f4, f8, c2, c4, c8, f2, ud, ud, ud, f4},
                /* f4 */ {f4, f4, f4, f4, f4, f4, f4, f8, c4, c4, c8, f4, ud, ud, ud, f4},
                /* f8 */ {f8, f8, f8, f8, f8, f8, f8, f8, c8, c8, c8, f8, ud, ud, ud, f8},

And a new ‘reference’ implementation in primtorch which should match:

github.com

pytorch/pytorch/blob/master/torch/_prims/utils.py#L676

      
        
            class ELEMENTWISE_TYPE_PROMOTION_KIND(Enum):
                DEFAULT = (0,)
                INT_TO_FLOAT = (1,)
                ALWAYS_BOOL = (2,)
                OP_MATH = (3,)
                COMPLEX_TO_FLOAT = (4,)
                BOOL_TO_LONG = (5,)
            
            

            
# TODO: document type promotion kinds
            def elementwise_dtypes(
                *_args, type_promotion_kind: ELEMENTWISE_TYPE_PROMOTION_KIND
            ) -> Tuple[torch.dtype, torch.dtype]:
                """
                Computes the computation and result dtypes for elementwise type promotion
                on the given arguments and with the given elementwise type promotion kind.
            
            
    Note that not all inputs to an elementwise operation necessarily participate in type promotion.
                For example, the "alpha" parameter of torch.add does not participate in type promotion,
                although it is cast to the Python type corresponding to the computation dtype that
                the type promotion algorithm determines.

Does this help?