Add activation functions (ReLU and SiLU for now) for structured spars…

…e linear operator (#101339) Pull Request resolved: #101339 Approved by: https://github.com/cpuhrsch
pytorch · May 22, 2023 · f9a327a · f9a327a
1 parent 320ebb4
commit f9a327a
Show file tree

Hide file tree

Showing 3 changed files with 370 additions and 118 deletions.
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -3210,7 +3210,7 @@
     MkldnnCPU: mkldnn_linear_backward
   autogen: mkldnn_linear_backward.out
 
-- func: _structured_sparse_linear(Tensor input, Tensor weight, Tensor mask_or_meta, *, Tensor? bias=None) -> (Tensor, Tensor)
+- func: _structured_sparse_linear(Tensor input, Tensor weight, Tensor mask_or_meta, *, Tensor? bias=None, str? activation=None) -> (Tensor, Tensor)
   dispatch:
     CUDA: _structured_sparse_linear