content: # Copyright (c) Facebook, Inc. and its affiliates.
  
  import torch
  import torch.nn as nn
  
  
  class Mlp(nn.Module):
      def __init__(
          self,
          in_features,
          hidden_features=None,
          out_features=None,
          act_layer=nn.GELU,
          drop_rate=0.0,
      ):
          super().__init__()
          self.drop_rate = drop_rate
          out_features = out_features or in_features
          hidden_features = hidden_features or in_features
          self.fc1 = nn.Linear(in_features, hidden_features)
          self.act = act_layer()
          self.fc2 = nn.Linear(hidden_features, out_features)
          if self.drop_rate > 0.0:
              self.drop = nn.Dropout(drop_rate)
  
      def forward(self, x):
          x = self.fc1(x)
          x = self.act(x)
          if self.drop_rate > 0.0:
              x = self.drop(x)
          x = self.fc2(x)
          if self.drop_rate > 0.0:
              x = self.drop(x)
          return x
  
  
  class Permute(nn.Module):
      def __init__(self, dims):
          super().__init__()
          self.dims = dims
  
      def forward(self, x):
          return x.permute(*self.dims)
  
  
  def drop_path(x, drop_prob: float = 0.0, training: bool = False):
      """
      Stochastic Depth per sample.
      """
      if drop_prob == 0.0 or not training:
          return x
      keep_prob = 1 - drop_prob
      shape = (x.shape[0],) + (1,) * (
          x.ndim - 1
      )  # work with diff dim tensors, not just 2D ConvNets
      mask = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
      mask.floor_()  # binarize
      output = x.div(keep_prob) * mask
      return output
  
  
  class DropPath(nn.Module):
      """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""
  
      def __init__(self, drop_prob=None):
          super(DropPath, self).__init__()
          self.drop_prob = drop_prob
  
      def forward(self, x):
          return drop_path(x, self.drop_prob, self.training)
  