[TORCH] Add support for aten.heaviside Op #4220
Conversation
|
@stellaraccident @vivekkhandelwal1 @penguin-wwy @zjgarvey @AmosLewis, I’d be grateful if any of you could take a look at this PR. Your feedback would be greatly appreciated! |
| class ElementwiseHeavisideIntModule(torch.nn.Module): | ||
| def __init__(self): | ||
| super().__init__() | ||
|
|
||
| @export | ||
| @annotate_args([None, ([-1, -1], torch.int32, True), ([-1], torch.int32, True)]) | ||
| def forward(self, x, values): | ||
| return torch.heaviside(x, values) | ||
|
|
||
|
|
||
| @register_test_case(module_factory=lambda: ElementwiseHeavisideIntModule()) | ||
| def ElementwiseHeavisideIntModule_basic(module, tu: TestUtils): | ||
| module.forward( | ||
| tu.randint(5, 1, low=-100, high=1000).to(torch.int32), | ||
| tu.randint(1, low=-100, high=1000).to(torch.int32), | ||
| ) |
There was a problem hiding this comment.
Why did you cast to int32 here?
| SmallVector<int64_t> broadcastShape; | ||
| SmallVector<Value> broadcastShapeValue; | ||
| computeBroadcastShape(rewriter, loc, input, value, broadcastShape, | ||
| broadcastShapeValue); | ||
|
|
||
| auto broadcastType = ValueTensorType::get( | ||
| op.getContext(), llvm::ArrayRef(broadcastShape), resultTy.getDtype()); | ||
| auto boolBroadcastType = ValueTensorType::get( | ||
| op.getContext(), llvm::ArrayRef(broadcastShape), rewriter.getI1Type()); | ||
| Value indexBroadcastShapeTorchList = rewriter.create<PrimListConstructOp>( | ||
| loc, Torch::ListType::get(Torch::IntType::get(op.getContext())), | ||
| broadcastShapeValue); | ||
| auto inputBroadcasted = rewriter.create<AtenBroadcastToOp>( | ||
| loc, broadcastType, input, indexBroadcastShapeTorchList); | ||
| auto valueBroadcasted = rewriter.create<AtenBroadcastToOp>( | ||
| loc, broadcastType, value, indexBroadcastShapeTorchList); |
There was a problem hiding this comment.
I think this is not needed. Since you are decomposing this op into elementwise ops, the broadcasting part will be handled during Torch->Linalg lowering.
| // Compute mask: isnan(input) | ||
| auto isNan = | ||
| rewriter.create<AtenIsnanOp>(loc, boolBroadcastType, inputBroadcasted); |
There was a problem hiding this comment.
I did not see the mention of this case here: https://docs.pytorch.org/docs/stable/generated/torch.heaviside.html. Can you please share any reference?
There was a problem hiding this comment.
Thanks for the review, @vivekkhandelwal1.
I tested this behavior with PyTorch — if the input contains NaN values, they are replaced with 0.
To handle this explicitly, I used AtenIsnanOp to detect NaN values
input=torch.tensor([[0,float('nan')]])
values=torch.tensor([2],dtype=torch.float32)
torch.heaviside(input,values)
Output
tensor([[2., 0.]])
Ref: https://github.com/pytorch/pytorch/blob/main/torch/_refs/__init__.py#L1448
This implementation addresses and closes #4211