class_docstr: Applies a stack of Perceptron modules sequentially (i.e. Multi-Layer Perceptron).
  
  Args:
      in_size (int): `in_size` of the input.
      layer_sizes (List[int]): `out_size` of each Perceptron module.
      bias (bool): if set to False, the layer will not learn an additive bias.
          Default: True.
      activation (str, Union[Callable[[], torch.nn.Module], torch.nn.Module, Callable[[torch.Tensor], torch.Tensor]]):
          the activation function to apply to the output of linear transformation of
          each Perceptron module.
          If `activation` is a `str`, we currently only support the follow strings, as
          "relu", "sigmoid", and "swish_layernorm".
          If `activation` is a `Callable[[], torch.nn.Module]`, `activation()` will be
          called once per Perceptron module to generate the activation module for that
          Perceptron module, and the parameters won't be shared between those activation
          modules.
          One use case is when all the activation modules share the same constructor
          arguments, but don't share the actual module parameters.
          Default: torch.relu.
      device (Optional[torch.device]): default compute device.
  
  Example::
  
      batch_size = 3
      in_size = 40
      input = torch.randn(batch_size, in_size)
  
      layer_sizes = [16, 8, 4]
      mlp_module = MLP(in_size, layer_sizes, bias=True)
      output = mlp_module(input)
      assert list(output.shape) == [batch_size, layer_sizes[-1]]
class_docstr_tok: ["Applies","a","stack","of","Perceptron","modules","sequentially","(","i.e",".","Multi-Layer","Perceptron",")",".","Args",":","in_size","(","int",")",":","`","in_size","`","of","the","input",".","layer_sizes","(","List","[","int","]",")",":","`","out_size","`","of","each","Perceptron","module",".","bias","(","bool",")",":","if","set","to","False",",","the","layer","will","not","learn","an","additive","bias",".","Default",":","True",".","activation","(","str",",","Union","[","Callable","[","[","]",",","torch.nn.Module","]",",","torch.nn.Module",",","Callable","[","[","torch.Tensor","]",",","torch.Tensor","]","]",")",":","the","activation","function","to","apply","to","the","output","of","linear","transformation","of","each","Perceptron","module",".","If","`","activation","`","is","a","`","str","`",",","we","currently","only","support","the","follow","strings",",","as","``","relu","''",",","``","sigmoid","''",",","and","``","swish_layernorm","''",".","If","`","activation","`","is","a","`","Callable","[","[","]",",","torch.nn.Module","]","`",",","`","activation","(",")","`","will","be","called","once","per","Perceptron","module","to","generate","the","activation","module","for","that","Perceptron","module",",","and","the","parameters","wo","n't","be","shared","between","those","activation","modules",".","One","use","case","is","when","all","the","activation","modules","share","the","same","constructor","arguments",",","but","do","n't","share","the","actual","module","parameters",".","Default",":","torch.relu",".","device","(","Optional","[","torch.device","]",")",":","default","compute","device",".","Example",":",":","batch_size","=","3","in_size","=","40","input","=","torch.randn","(","batch_size",",","in_size",")","layer_sizes","=","[","16",",","8",",","4","]","mlp_module","=","MLP","(","in_size",",","layer_sizes",",","bias=True",")","output","=","mlp_module","(","input",")","assert","list","(","output.shape",")","==","[","batch_size",",","layer_sizes","[","-1","]","]"]
code: def __init__(
      self,
      in_size: int,
      layer_sizes: List[int],
      bias: bool = True,
      activation: Union[
          str,
          Callable[[], torch.nn.Module],
          torch.nn.Module,
          Callable[[torch.Tensor], torch.Tensor],
      ] = torch.relu,
      device: Optional[torch.device] = None,
      dtype: torch.dtype = torch.float32,
  ) -> None:
      super().__init__()
  
      if activation == "relu":
          activation = torch.relu
      elif activation == "sigmoid":
          activation = torch.sigmoid
  
      if not isinstance(activation, str):
          self._mlp: torch.nn.Module = torch.nn.Sequential(
              *[
                  Perceptron(
                      layer_sizes[i - 1] if i > 0 else in_size,
                      layer_sizes[i],
                      bias=bias,
                      activation=extract_module_or_tensor_callable(
                          activation
                      ),
                      device=device,
                      dtype=dtype,
                  )
                  for i in range(len(layer_sizes))
              ]
          )
      else:
          if activation == "swish_layernorm":
              self._mlp: torch.nn.Module = torch.nn.Sequential(
                  *[
                      Perceptron(
                          layer_sizes[i - 1] if i > 0 else in_size,
                          layer_sizes[i],
                          bias=bias,
                          activation=SwishLayerNorm(
                              layer_sizes[i], device=device
                          ),
                          device=device,
                      )
                      for i in range(len(layer_sizes))
                  ]
              )
          else:
              assert (
                  ValueError
              ), "This MLP only support str version activation function of relu, sigmoid, and swish_layernorm"
  
file_path: torchrec/torchrec/modules/mlp.py
imports: from typing import Callable, List, Optional, Union
  import torch
  from torch import nn
  from torchrec.modules.activation import SwishLayerNorm
  from torchrec.modules.utils import extract_module_or_tensor_callable
label_desc: Use image node_id 1 to create a new MLP object from inherited base classes: torch.nn with example: obj = MLP(in_size, layer_sizes, bias, activation, device, dtype)