OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 37 Packages Projects Releases Wiki Activity

[MPS/inductor] Add support for modified_scaled_bessel_k{0,1} (#149794)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/149794
Approved by: https://github.com/malfet
This commit is contained in:
Davide Italiano 2025-03-22 15:41:40 +00:00 committed by PyTorch MergeBot
parent 6bbe8dbd63
commit 2b848ab192
2 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
test/inductor/test_mps_basic.py
View File

@ -99,6 +99,8 @@ class MPSBasicTests(TestCase):
"modified_bessel_i1", "modified_bessel_i1",
"modified_bessel_k0", "modified_bessel_k0",
"modified_bessel_k1", "modified_bessel_k1",
"scaled_modified_bessel_k0",
"scaled_modified_bessel_k1",
"entr", "entr",
] ]

8
torch/_inductor/codegen/mps.py
View File

@ -418,6 +418,14 @@ class MetalOverrides(OpOverrides):
def modified_bessel_k1(x: CSEVariable) -> str: def modified_bessel_k1(x: CSEVariable) -> str:
return f"c10::metal::modified_bessel_k1_forward({x})" return f"c10::metal::modified_bessel_k1_forward({x})"
@staticmethod
def scaled_modified_bessel_k0(x: CSEVariable) -> str:
return f"c10::metal::scaled_modified_bessel_k0_forward({x})"
@staticmethod
def scaled_modified_bessel_k1(x: CSEVariable) -> str:
return f"c10::metal::scaled_modified_bessel_k1_forward({x})"
MetalOverrides._initialize_pointwise_overrides("mps") MetalOverrides._initialize_pointwise_overrides("mps")