Fix: #7084

This PR fixes a data-type related problem for mul operation. It does so by creating a structure OpConfig that behaves similarly to DoBinaryOp. The difference is that it takes care of pre/post-processing of inputs and outputs, casting them to the correct data-type.

Problem

t = torch.rand(10, dtype=torch.half).to(xm.xla_device())
s = torch.tensor(10, dtype=torch.double).to(xm.xla_device())
out = torch.mul_(t, s)

• Tensor.mul_ is dispatched to its CompositeExplicitAutograd kernel
  • It wraps the scalar into a tensor, and calls torch.mul (functional version)
• DoBinaryOp is called
  • Computes at::result_type (let's call it common_dtype) and passes it on to bin_op
  • Note that UnwrapNumber does nothing, since s is a tensor with is_wrapped_number_ unset
• The computed common_dtype is passed on to tensor_methods::mul
  • Creates an IR node with data-type common_dtype
  • Does nothing with its inputs
• Later, when BuildMul is called, we have 2 XlaOp with different data-types: f16 and f64
  • BuildMul promotes f16 to f64
• The output is common_dtype (torch.float16), but the actal XlaOp is f64

Solution

Following PyTorch behavior [1, 2, 3], I created OpConfig: a structure that let us specify common pre/post-processing on inputs and outputs.

Affected Models

• timm_nfnet (training+nondynamo)

cc @miladm @JackCaoG @lezcano