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

…e linear operator ghstack-source-id: 1114550774c7f46e7b204f0137ab7599bfec1526 Pull Request resolved: #101339
pytorch · May 13, 2023 · 8be3c1e · 8be3c1e
1 parent 317a791
commit 8be3c1e
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