Update

manuelcandales · manuelcandales · commit e0ae6be4681e · 2025-06-11T16:49:56.000-04:00
[ghstack-poisoned]
diff --git a/aten/src/ATen/native/mps/kernels/ActivationKernel.metal b/aten/src/ATen/native/mps/kernels/ActivationKernel.metal
@@ -11,7 +11,20 @@ struct hardshrink_functor {
   }
 };
 
-struct hardshrink_backward_functor {
+struct softshrink_functor {
+  template <typename T>
+  inline T operator()(const T x, const T lambda) {
+    if (x > lambda) {
+      return x - lambda;
+    } else if (x < -lambda) {
+      return x + lambda;
+    } else {
+      return T(0);
+    }
+  }
+};
+
+struct shrink_backward_functor {
   template <typename T>
   inline T operator()(const T grad_output, const T x, const T lambda) {
     return (x >= -lambda && x <= lambda) ? T(0) : grad_output;
@@ -24,10 +37,16 @@ REGISTER_UNARY_ALPHA_OP(hardshrink, half, half, half);
 REGISTER_UNARY_ALPHA_OP(hardshrink, bfloat, bfloat, bfloat);
 #endif
 
-REGISTER_BINARY_ALPHA_OP(hardshrink_backward, float, float, float);
-REGISTER_BINARY_ALPHA_OP(hardshrink_backward, half, half, half);
+REGISTER_UNARY_ALPHA_OP(softshrink, float, float, float);
+REGISTER_UNARY_ALPHA_OP(softshrink, half, half, half);
+#if __METAL_VERSION__ >= 310
+REGISTER_UNARY_ALPHA_OP(softshrink, bfloat, bfloat, bfloat);
+#endif
+
+REGISTER_BINARY_ALPHA_OP(shrink_backward, float, float, float);
+REGISTER_BINARY_ALPHA_OP(shrink_backward, half, half, half);
 #if __METAL_VERSION__ >= 310
-REGISTER_BINARY_ALPHA_OP(hardshrink_backward, bfloat, bfloat, bfloat);
+REGISTER_BINARY_ALPHA_OP(shrink_backward, bfloat, bfloat, bfloat);
 #endif
 
 struct hardsigmoid_functor {
@@ -128,39 +147,3 @@ REGISTER_BINARY_ALPHA_OP(leaky_relu_backward, half, half, half);
 #if __METAL_VERSION__ >= 310
 REGISTER_BINARY_ALPHA_OP(leaky_relu_backward, bfloat, bfloat, bfloat);
 #endif
-
-struct softshrink_functor {
-  template <typename T>
-  inline T operator()(const T x, const T lambda) {
-    if (x > lambda) {
-      return x - lambda;
-    } else if (x < -lambda) {
-      return x + lambda;
-    } else {
-      return T(0);
-    }
-  }
-};
-
-struct softshrink_backward_functor {
-  template <typename T>
-  inline T operator()(const T grad_output, const T self, const T lambda) {
-    if (self > lambda || self < -lambda) {
-      return grad_output;
-    } else {
-      return T(0);
-    }
-  }
-};
-
-REGISTER_UNARY_ALPHA_OP(softshrink, float, float, float);
-REGISTER_UNARY_ALPHA_OP(softshrink, half, half, half);
-#if __METAL_VERSION__ >= 310
-REGISTER_UNARY_ALPHA_OP(softshrink, bfloat, bfloat, bfloat);
-#endif
-
-REGISTER_BINARY_ALPHA_OP(softshrink_backward, float, float, float);
-REGISTER_BINARY_ALPHA_OP(softshrink_backward, half, half, half);
-#if __METAL_VERSION__ >= 310
-REGISTER_BINARY_ALPHA_OP(softshrink_backward, bfloat, bfloat, bfloat);
-#endif
diff --git a/aten/src/ATen/native/mps/operations/ActivationKernel.mm b/aten/src/ATen/native/mps/operations/ActivationKernel.mm
@@ -17,8 +17,12 @@ static void hardshrink_kernel(TensorIteratorBase& iter, const Scalar& lambda = 0
   lib.exec_unary_kernel(iter, "hardshrink", lambda);
 }
 
+static void softshrink_kernel(TensorIteratorBase& iter, const Scalar& lambda = 0.5) {
+  lib.exec_unary_kernel(iter, "softshrink", lambda);
+}
+
 static void shrink_backward_kernel(TensorIteratorBase& iter, const Scalar& lambda = 0.5) {
-  lib.exec_binary_kernel(iter, "hardshrink_backward", lambda);
+  lib.exec_binary_kernel(iter, "shrink_backward", lambda);
 }
 
 static void hardsigmoid_kernel(TensorIteratorBase& iter) {
@@ -45,18 +49,14 @@ static void leaky_relu_backward_kernel(TensorIteratorBase& iter, const Scalar& n
   lib.exec_binary_kernel(iter, "leaky_relu_backward", negative_slope);
 }
 
-static void softshrink_kernel(TensorIteratorBase& iter, const Scalar& lambda = 0.5) {
-  lib.exec_unary_kernel(iter, "softshrink", lambda);
-}
-
 REGISTER_DISPATCH(hardshrink_stub, hardshrink_kernel);
+REGISTER_DISPATCH(softshrink_stub, softshrink_kernel);
 REGISTER_DISPATCH(shrink_backward_stub, shrink_backward_kernel);
 REGISTER_DISPATCH(hardsigmoid_stub, hardsigmoid_kernel);
 REGISTER_DISPATCH(hardsigmoid_backward_stub, hardsigmoid_backward_kernel);
 REGISTER_DISPATCH(hardswish_stub, hardswish_kernel);
 REGISTER_DISPATCH(hardswish_backward_stub, hardswish_backward_kernel);
 REGISTER_DISPATCH(leaky_relu_stub, leaky_relu_kernel);
 REGISTER_DISPATCH(leaky_relu_backward_stub, leaky_relu_backward_kernel);
-REGISTER_DISPATCH(softshrink_stub, softshrink_kernel);
 
 } // namespace at::native
