saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-05-29 23:06:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add aten::_softmax to eager ops. (#11820)

Browse source
This commit is contained in:
Wil Brady 2022-06-13 13:05:26 -04:00 committed by GitHub
parent 7582644f57
commit b0e027c661
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 9 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
orttraining/orttraining/eager/opgen/opgen/atenops.py
View file

@ -138,6 +138,7 @@ hand_implemented = {
"aten::_local_scalar_dense": MakeTorchFallback(),
"aten::gt.Scalar_out": MakeTorchFallback(),
"aten::equal": MakeTorchFallback(),
"aten::_softmax": Softmax("self", axis="dim"),
}
# Signature of gelu_backward was changed in this commit id 983ba5e585485ed61a0c0012ef6944f5685e3d97 and PR 61439

8
orttraining/orttraining/eager/test/ort_ops.py
View file

@ -144,6 +144,14 @@ class OrtOpTests(unittest.TestCase):
cpu_tensor_copied = ort_tensor.cpu()
assert cpu_tensor_copied.stride() == (0, 0, 0)
def test_softmax(self):
device = self.get_device()
cpu_tensor = torch.rand(3, 5)
ort_tensor = cpu_tensor.to(device)
cpu_result = torch.softmax(cpu_tensor, dim=1)
ort_result = torch.softmax(ort_tensor, dim=1)
assert torch.allclose(cpu_result, ort_result.cpu())
if __name__ == "__main__":
unittest.main()