class_docstr: An implementation of a simple dense feedforward network taking a context variable, for use in, e.g.,
  some conditional flows such as :class:`pyro.distributions.transforms.ConditionalAffineCoupling`.
  
  Example usage:
  
  >>> input_dim = 10
  >>> context_dim = 5
  >>> x = torch.rand(100, input_dim)
  >>> z = torch.rand(100, context_dim)
  >>> nn = ConditionalDenseNN(input_dim, context_dim, [50], param_dims=[1, input_dim, input_dim])
  >>> a, b, c = nn(x, context=z)  # parameters of size (100, 1), (100, 10), (100, 10)
  
  :param input_dim: the dimensionality of the input
  :type input_dim: int
  :param context_dim: the dimensionality of the context variable
  :type context_dim: int
  :param hidden_dims: the dimensionality of the hidden units per layer
  :type hidden_dims: list[int]
  :param param_dims: shape the output into parameters of dimension (p_n,) for p_n in param_dims
      when p_n > 1 and dimension () when p_n == 1. The default is [1, 1], i.e. output two parameters of dimension ().
  :type param_dims: list[int]
  :param nonlinearity: The nonlinearity to use in the feedforward network such as torch.nn.ReLU(). Note that no
      nonlinearity is applied to the final network output, so the output is an unbounded real number.
  :type nonlinearity: torch.nn.Module
code: def __init__(
      self,
      input_dim,
      context_dim,
      hidden_dims,
      param_dims=[1, 1],
      nonlinearity=torch.nn.ReLU(),
  ):
      super().__init__()
  
      self.input_dim = input_dim
      self.context_dim = context_dim
      self.hidden_dims = hidden_dims
      self.param_dims = param_dims
      self.count_params = len(param_dims)
      self.output_multiplier = sum(param_dims)
  
      # Calculate the indices on the output corresponding to each parameter
      ends = torch.cumsum(torch.tensor(param_dims), dim=0)
      starts = torch.cat((torch.zeros(1).type_as(ends), ends[:-1]))
      self.param_slices = [
          slice(s.item(), e.item()) for s, e in zip(starts, ends)
      ]
  
      # Create masked layers
      layers = [
          torch.nn.Linear(input_dim + context_dim, hidden_dims[0])
      ]
      for i in range(1, len(hidden_dims)):
          layers.append(
              torch.nn.Linear(hidden_dims[i - 1], hidden_dims[i])
          )
      layers.append(
          torch.nn.Linear(hidden_dims[-1], self.output_multiplier)
      )
      self.layers = torch.nn.ModuleList(layers)
  
      # Save the nonlinearity
      self.f = nonlinearity
  
file_path: pyro/pyro/nn/dense_nn.py
imports: import torch
label_desc: Use image node_id 1 to create a new ConditionalDenseNN object from inherited base classes: torch.nn with example: obj = ConditionalDenseNN(input_dim, context_dim, hidden_dims, param_dims, nonlinearity)