Is the torchax test failure expected?

While I understand the suggestion from #8545 to make this backend-specific, I believe a full revert is more appropriate atleast for release branch r2.9:

1. The change affects the fundamental numerical behavior of torch.mul, which is called extensively throughout any model. Even small precision differences compound in:
   • Attention mechanisms with large sequence lengths
   • Gradient accumulation/reduction over many steps and likely XLA based compilers need to explicitly handle this scenario.

2. Risk vs. benefit:
   • Risk: Our testing shows concrete accuracy regressions and are blocking for production workloads(mixtral, llama) on current compiler, which I guess might impact other users too.
   • Benefit: The original issue [Q][GPU][BF16] torch.mul is lowered to HLO as an f32 multiply #8545 was a question about the seeming unnecessary upcast/downcast that also appears in PyTorch CUDA. The upcast/downcast can be already removed by using PyTorch autocast so there's no need to fix [Q][GPU][BF16] torch.mul is lowered to HLO as an f32 multiply #8545 for r2.9. We can keep this in master branch for further investigation.

Let me know your thoughts?

Thank you for your detailed analysis.

I'm sorry, but I still think that the best way to go about this is to make it backend specific. Actually, now that I'm thinking about it, I think it would make more sense to do it in the lowering step.

That said, since we don't really want to introduce a lot of changes here, and the infrastructure is already in place, I would say that we should make that use_opmathtype_for_compute() call backend specific.

It's a simple change (something like this condition for TPU) that fixes this problem for the Neuron backend, while leaving the other backends with the "backend-independent" choice of not doing that.

2. The original issue [Q][GPU][BF16] torch.mul is lowered to HLO as an f32 multiply #8545 was a question about the seeming unnecessary upcast/downcast that also appears in PyTorch CUDA. The upcast/downcast can be already removed by using PyTorch autocast so there's no need to fix [Q][GPU][BF16] torch.mul is lowered to HLO as an f32 multiply #8545 for r2.9. We can keep this in master branch for further investigation.

That "seeming unnecessary upcast/downcast" is exactly what we are talking about, here. And, no, PyTorch autocast does not solve this issue (see this comment).

Could you rebase this PR, so that there are only your commits?

Could you rebase this PR, so that there are only your commits?

updated

@ysiraichi CI is completed.

Thanks @rajkthakur @ysiraichi . It is merged now and ready to go. @bhavya01 will make a final 2.9 release candidate.