[MPS] Add inductor support for spherical_bessel_j0. (#147650)

dcci · pytorchmergebot · commit 6a5e3917a74d · 2025-02-23T00:32:36.000Z
Counterpart to my previous patch that added support for the op in eager. Pull Request resolved: #147650 Approved by: https://github.com/jansel
diff --git a/test/inductor/test_mps_basic.py b/test/inductor/test_mps_basic.py
@@ -114,6 +114,11 @@ def test_pointwise_zeta(self):
             check_lowp=False,
         )
 
+    def test_pointwise_spherical_bessel_j0(self):
+        self.common(
+            torch.special.spherical_bessel_j0, (torch.rand(128, 128),), check_lowp=False
+        )
+
     def test_broadcast(self):
         self.common(torch.add, (torch.rand(32, 1024), torch.rand(1024)))
 
diff --git a/torch/_inductor/codegen/mps.py b/torch/_inductor/codegen/mps.py
@@ -363,6 +363,10 @@ def pow(a: CSEVariable, b: CSEVariable) -> str:
     def zeta(a: CSEVariable, b: CSEVariable) -> str:
         return f"c10::metal::zeta({a}, {b})"
 
+    @staticmethod
+    def spherical_bessel_j0(x: CSEVariable) -> str:
+        return f"c10::metal::spherical_bessel_j0({x})"
+
 
 MetalOverrides._initialize_pointwise_overrides("mps")
 
