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

…e linear operator ghstack-source-id: dbccf81feeae952f008d29b90ff12b3e2fa617ca Pull Request resolved: #101339
pytorch · May 13, 2023 · 93e30f5 · 93e30f5
1 parent 317a791
commit 93e30f5
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
@@ -3191,7 +3191,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