separate backward function to not repeat with hardshrink backward

igm503 · igm503 · commit 5569a5c2e8c1 · 2023-10-08T03:00:07.000-05:00
diff --git a/aten/src/ATen/native/mps/operations/Activation.mm b/aten/src/ATen/native/mps/operations/Activation.mm
@@ -1524,8 +1524,11 @@ Tensor mish_backward_mps(const Tensor& grad_output, const Tensor& self) {
   }
 }
 
-TORCH_IMPL_FUNC(softshrink_backward_out_mps)
-(const Tensor& grad_output, const Tensor& self, const Scalar& lambd, const Tensor& grad_input) {
+static void shrink_backward_out_mps(const Tensor& grad_output,
+                                    const Tensor& self,
+                                    const Scalar& lambd,
+                                    const Tensor& grad_input,
+                                    std::string op_name) {
   using namespace mps;
   TORCH_CHECK(self.is_mps());
 
@@ -1545,8 +1548,7 @@ Tensor mish_backward_mps(const Tensor& grad_output, const Tensor& self) {
   MPSStream* stream = getCurrentMPSStream();
 
   @autoreleasepool {
-    string key =
-        "softshrink_backward_out_mps:" + getTensorsStringKey({self}) + ":" + std::to_string(lambd.to<double>());
+    string key = op_name + ":" + getTensorsStringKey({self}) + ":" + std::to_string(lambd.to<double>());
 
     auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
       MPSGraphTensor* gradOutputTensor = mpsGraphRankedPlaceHolder(mpsGraph, grad_output);
@@ -1588,9 +1590,15 @@ Tensor mish_backward_mps(const Tensor& grad_output, const Tensor& self) {
     NSDictionary<MPSGraphTensor*, MPSGraphTensorData*>* results =
         @{gradInputPlaceholder.getMPSGraphTensor() : gradInputPlaceholder.getMPSGraphTensorData()};
     runMPSGraph(stream, cachedGraph->graph(), feeds, results);
+    return;
   }
 }
 
+TORCH_IMPL_FUNC(softshrink_backward_out_mps)
+(const Tensor& grad_output, const Tensor& self, const Scalar& lambd, const Tensor& grad_input) {
+  return shrink_backward_out_mps(grad_output, self, lambd, grad_input, "softshrink_backward_out_mps");
+}
+
 Tensor prelu_mps(const Tensor& self, const Tensor& weight_) {
   using namespace mps;
 
