docstr: Forward pass.
  
  X: (batch, dim, lengths...)
docstr_tok: ["Forward","pass",".","X",":","(","batch",",","dim",",","lengths","...",")"]
code: def forward(self, X):
      """Forward pass.
  
      X: (batch, dim, lengths...)
      """
      if self.training:
          if not self.transposed:
              X = rearrange(X, "b d ... -> b ... d")
          # binomial = torch.distributions.binomial.Binomial(
          #   probs=1-self.p) # This is incredibly slow
          mask_shape = (
              X.shape[:2] + (1,) * (X.ndim - 2) if self.tie else X.shape
          )
          # mask = self.binomial.sample(mask_shape)
          mask = torch.rand(*mask_shape, device=X.device) < 1.0 - self.p
          X = X * mask * (1.0 / (1 - self.p))
          if not self.transposed:
              X = rearrange(X, "b ... d -> b d ...")
          return X
      return X
  
imports: import math
  from functools import partial
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from opt_einsum import contract
label_desc: Use image node_id 2 for calling the DropoutNd obj's underlying member method code with example usage: obj.forward(X) and returns: X, X