Add inplace parameter for relu

Author: stever178

src/ntops/kernels/relu.py

@@ -6,10 +6,22 @@
 from ntops.kernels.element_wise import arrangement
 
 
-def application(input, output):
+def default_application(input, output):
     output = max(0.0, input)  # noqa: F841
 
 
+def True_application(input, output):
+    input = max(0.0, input)  
+    output = input  # noqa: F841
+
+
 @functools.cache
-def make(ndim):
-    return ninetoothed.make(arrangement, application,(Tensor(ndim), Tensor(ndim)))
+def make(ndim, replace):
+    if replace == True:
+        application = True_application
+    else:
+        application = default_application
+
+    tensors = (Tensor(ndim), Tensor(ndim))
+
+    return ninetoothed.make(arrangement, application, tensors)

src/ntops/torch.py

@@ -120,15 +120,14 @@ def mul(input, other, *, out=None):
     return out
 
 
-def relu(input, out=None):
-    if out is None:
-        out = torch.empty_like(input)
+def relu(input, inplace=False):
+    output = torch.empty_like(input)
 
-    kernel = ntops.kernels.relu.make(input.ndim)
+    kernel = ntops.kernels.relu.make(input.ndim, inplace)
 
-    kernel(input, out)
+    kernel(input, output)
 
-    return out
+    return output
 
 
 def rsqrt(input, *, out=None):

tests/test_relu.py

@@ -14,7 +14,10 @@ def test_cuda(shape, dtype, atol, rtol):
 
     input = torch.randn(shape, dtype=dtype, device=device)
 
-    ninetoothed_output = ntops.torch.relu(input)
-    reference_output = F.relu(input)
+    for replace in (False, True):
+        ninetoothed_output = ntops.torch.relu(input, replace)
+        reference_output = F.relu(input, replace)
 
-    assert torch.allclose(ninetoothed_output, reference_output, atol=atol, rtol=rtol)
+        assert torch.allclose(
+            ninetoothed_output, reference_output, atol=atol, rtol=rtol
+        )