code: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.distributions as td
  
  class Flow(nn.Module):
      """
      Building both normalizing flows and neural flows.
  
      Example:
      >>> import stribor as st
      >>> torch.manual_seed(123)
      >>> dim = 2
      >>> flow = st.Flow(st.UnitNormal(dim), [st.Affine(dim)])
      >>> x = torch.rand(1, dim)
      >>> y, ljd = flow(x)
      >>> y_inv, ljd_inv = flow.inverse(y)
  
      Args:
          base_dist (Type[torch.distributions]): Base distribution
          transforms (List[st.flows]): List of invertible transformations
      """
      def __init__(self, base_dist=None, transforms=[]):
          super().__init__()
          self.base_dist = base_dist
          self.transforms = nn.ModuleList(transforms)
  
      def forward(self, x, latent=None, mask=None, t=None, reverse=False, **kwargs):
          """
          Args:
              x (tensor): Input sampled from base density with shape (..., dim)
              latent (tensor, optional): Conditional vector with shape (..., latent_dim)
                  Default: None
              mask (tensor): Masking tensor with shape (..., 1)
                  Default: None
              t (tensor, optional): Flow time end point. Default: None
              reverse (bool, optional): Whether to perform an inverse. Default: False
  
          Returns:
              y (tensor): Output that follows target density (..., dim)
              log_jac_diag (tensor): Log-Jacobian diagonal (..., dim)
          """
          transforms = self.transforms[::-1] if reverse else self.transforms
          _mask = 1 if mask is None else mask
  
          log_jac_diag = torch.zeros_like(x).to(x)
          for f in transforms:
              if reverse:
                  x, ld = f.inverse(x * _mask, latent=latent, mask=mask, t=t, **kwargs)
              else:
                  x, ld = f.forward(x * _mask, latent=latent, mask=mask, t=t, **kwargs)
              log_jac_diag += ld * _mask
          return x, log_jac_diag
  
      def inverse(self, y, latent=None, mask=None, t=None, **kwargs):
          """ Inverse of forward function with the same arguments. """
          return self.forward(y, latent=latent, mask=mask, t=t, reverse=True, **kwargs)
  
      def log_prob(self, x, **kwargs):
          """
          Calculates log-probability of a sample.
  
          Args:
              x (tensor): Input with shape (..., dim)
  
          Returns:
              log_prob (tensor): Log-probability of the input with shape (..., 1)
          """
          if self.base_dist is None:
              raise ValueError('Please define `base_dist` if you need log-probability')
          x, log_jac_diag = self.inverse(x, **kwargs)
          log_prob = self.base_dist.log_prob(x) + log_jac_diag.sum(-1)
          return log_prob.unsqueeze(-1)
  
      def sample(self, num_samples, latent=None, mask=None, **kwargs):
          """
          Transforms samples from the base to the target distribution.
          Uses reparametrization trick.
  
          Args:
              num_samples (tuple or int): Shape of samples
              latent (tensor): Latent conditioning vector with shape (..., latent_dim)
  
          Returns:
              x (tensor): Samples from target distribution with shape (*num_samples, dim)
          """
          if self.base_dist is None:
              raise ValueError('Please define `base_dist` if you need sampling')
          if isinstance(num_samples, int):
              num_samples = (num_samples,)
  
          x = self.base_dist.rsample(num_samples)
          x, log_jac_diag = self.forward(x, **kwargs)
          return x
  
apis: torch.zeros_like torch.nn.ModuleList