mirror of https://github.com/llvm/torch-mlir
e57d3f9774
- This commit adds E2E support for `aten.rand_like` and `aten.bernoulli_.Tensor` ops. - The `aten.bernoulli(x)` was implemented as: `aten.bernoulli(x) = rand_like(x) < 0.5`, assuming 0.5 as default probability, whereas according to the pytorch documentation: https://pytorch.org/docs/stable/generated/torch.bernoulli.html#torch.bernoulli the input x in `aten.bernoulli(x)` is itself a tensor containing probabilities to be used for drawing the binary random number. - So this commit fixes the `aten.bernoulli(x)` implementation as: `aten.bernoulli(x) = rand_like(x) < x`. - It also fixes the case where the input to `aten.bernoulli_.float` is an integer tensor. In this case the input must be casted to float type before passing it as operand to `aten.rand_like` op. `aten.bernoulli_.float(x, p) = rand_like(float(x)) < p`. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com> |
||
|---|---|---|
| .. | ||
| test | ||
| torch_mlir/dialects | ||
| torch_mlir_e2e_test | ||
| CMakeLists.txt | ||
| TorchMLIRModule.cpp | ||