Add softmax as another toy model

Summary: This is an interesting usecase to demonstrate on ARM Baremetal setup where they currently can't lower SoftMax. One of the only few from MV2. It is useful as a test as well while debugging this flow.

Differential Revision: D49899589

fbshipit-source-id: c0307c7452a69eb707429328ba6654b2b660ce23

Author: digantdesai

examples/models/__init__.py

@@ -11,6 +11,7 @@
     "linear": ("toy_model", "LinearModule"),
     "add": ("toy_model", "AddModule"),
     "add_mul": ("toy_model", "AddMulModule"),
+    "softmax" : ("toy_model", "SoftmaxModule"),
     "dl3": ("deeplab_v3", "DeepLabV3ResNet50Model"),
     "edsr": ("edsr", "EdsrModel"),
     "emformer_transcribe": ("emformer_rnnt", "EmformerRnntTranscriberModel"),

examples/models/toy_model/__init__.py

@@ -4,11 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .model import AddModule, AddMulModule, LinearModule, MulModule
+from .model import AddModule, AddMulModule, LinearModule, MulModule, SoftmaxModule
 
 __all__ = [
     AddModule,
     AddMulModule,
     LinearModule,
     MulModule,
+    SoftmaxModule,
 ]

examples/models/toy_model/model.py

@@ -75,3 +75,19 @@ def get_example_inputs(self):
     def get_compile_spec(self):
         max_value = self.get_example_inputs()[0].shape[0]
         return [CompileSpec("max_value", bytes([max_value]))]
+
+class SoftmaxModule(torch.nn.Module, EagerModelBase):
+    def __init__(self):
+        super().__init__()
+        self.softmax = torch.nn.Softmax()
+
+    def forward(self, x):
+        z = self.softmax(x)
+        return z
+
+    def get_eager_model(self) -> torch.nn.Module:
+        return self
+
+    def get_example_inputs(self):
+        return (torch.ones(2, 2),)
+