docs: def calc_nested_probs(nest_coefs,
                        index_coefs,
                        design,
                        rows_to_obs,
                        rows_to_nests,
                        chosen_row_to_obs=None,
                        return_type="long_probs",
                        *args,
                        **kwargs):
      """
      Parameters
      ----------
      nest_coefs : 1D or 2D ndarray.
          All elements should by ints, floats, or longs. If 1D, should have 1
          element for each nesting coefficient being estimated. If 2D, should
          have 1 column for each set of nesting coefficients being used to
          predict the probabilities of each alternative being chosen. There
          should be one row per nesting coefficient. Elements denote the inverse
          of the scale coefficients for each of the lower level nests.
      index_coefs : 1D or 2D ndarray.
          All elements should by ints, floats, or longs. If 1D, should have 1
          element for each utility coefficient being estimated (i.e.
          num_features). If 2D, should have 1 column for each set of coefficients
          being used to predict the probabilities of each alternative being
          chosen. There should be one row per index coefficient.
      design : 2D ndarray.
          There should be one row per observation per available alternative.
          There should be one column per utility coefficient being estimated. All
          elements should be ints, floats, or longs.
      rows_to_obs : 2D scipy sparse array.
          There should be one row per observation per available alternative and
          one column per observation. This matrix maps the rows of the design
          matrix to the unique observations (on the columns).
      rows_to_nests : 2D scipy sparse array.
          There should with one row per observation per available alternative and
          one column per nest. This matrix maps the rows of the design matrix to
          the unique nests (on the columns).
      chosen_row_to_obs : 2D scipy sparse array, or None, optional.
          There should be one row per observation per available alternative and
          one column per observation. This matrix indicates, for each observation
          (on the columns), which rows of the design matrix were the realized
          outcome. If an array is passed then an array of shape
          (num_observations,) can be returned and each element will be the
          probability of the realized outcome of the given observation.
          Default == None.
      return_type : str, optional.
          Indicates what object(s) are to be returned from the function. Valid
          values are: `['long_probs', 'chosen_probs', 'long_and_chosen_probs',
          'all_prob_dict']`. If `long_probs`, the long format probabilities (a 1D
          numpy array with one element per observation per available alternative)
          will be returned. If `chosen_probs`, a 1D numpy array with one element
          per observation will be returned, where the values are the
          probabilities of the chosen alternative for the given observation. If
          `long_and_chosen_probs`, a tuple of chosen_probs and long_probs will be
          returned. If `all_prob_dict`, a dictionary will be returned. The values
          will all be 1D numpy arrays of probabilities dictated by the value's
          corresponding key. The keys will be `long_probs`, `nest_choice_probs`,
          `prob_given_nest`, and `chosen_probs`. If chosen_row_to_obs is None,
          then `chosen_probs` will be None. If `chosen_row_to_obs` is passed,
          then `chosen_probs` will be a 1D array as described above.
          `nest_choice_probs` is of the same shape as `rows_to_nests` and it
          denotes the probability of each individual choosing each of the
          possible nests. `prob_given_nest` is of the same shape as `long_probs`
          and it denotes the probability of the individual associated with a
          given row choosing the alternative associated with that row, given that
          the individual chooses the nest that contains the given alternative.
          Default == `long_probs`.
  
      Returns
      -------
      See above for documentation of the `return_type` kwarg.
      """
      # Check for 2D index coefficients or nesting coefficients
      try:
          assert len(index_coefs.shape) <= 2
          assert (len(index_coefs.shape) == 1) or (index_coefs.shape[1] == 1)
          assert len(nest_coefs.shape) <= 2
          assert (len(nest_coefs.shape) == 1) or (nest_coefs.shape[1] == 1)
      except AssertionError:
          msg = "Support for 2D index_coefs or nest_coefs not yet implemented."
          raise NotImplementedError(msg)
  
      # Check for kwarg validity
      valid_return_types = ['long_probs',
                            'chosen_probs',
                            'long_and_chosen_probs',
                            'all_prob_dict']
      if return_type not in valid_return_types:
          msg = "return_type must be one of the following values: "
          raise ValueError(msg + str(valid_return_types))
  
      chosen_probs_needed = ['chosen_probs', 'long_and_chosen_probs']
      if chosen_row_to_obs is None and return_type in chosen_probs_needed:
          msg = "chosen_row_to_obs is None AND return_type in {}."
          raise ValueError(msg.format(chosen_probs_needed) +
                           "\nThis is invalid.")
  
      # Calculate the index for each alternative for each individual, V = X*beta
      index_vals = design.dot(index_coefs)
  
      # Get the long format nest parameters for each row of the design matrix
      long_nest_coefs = rows_to_nests.dot(nest_coefs)
  
      # Calculate the scaled index values (index / nest_param = V / lambda)
      scaled_index = index_vals / long_nest_coefs
  
      # Guard against overflow
      pos_inf_idx = np.isposinf(scaled_index)
      neg_inf_idx = np.isneginf(scaled_index)
      scaled_index[pos_inf_idx] = max_comp_value
      scaled_index[neg_inf_idx] = -1 * max_comp_value
  
      # Calculate the e^(scaled-index) = exp(V / lambda)
      exp_scaled_index = np.exp(scaled_index)
  
      # Guard against overflow
      inf_idx = np.isposinf(exp_scaled_index)
      exp_scaled_index[inf_idx] = max_comp_value
      # Guard against underflow. Note that I'm not sure this is the best place or
      # best way to perform such guarding. If all of an observations indices
      # suffer underflow, then we'll have 0 / 0 when calculating the
      # probabilities and I should use L'Hopital's rule to get the correct
      # probability. However, replacing underflowed values here may result in
      # incorrectly assigning probabilities of either zero for all alternatives
      # or 1 / num_alternatives for all alternatives.
      zero_idx = (exp_scaled_index == 0)
      exp_scaled_index[zero_idx] = min_comp_value
  
      # Calculate the log-sum for each nest, for each observation. Note that the
      # "*" is used to compute the dot product between the mapping matrix which
      # is a scipy.sparse matrix and the second term which is a scipy sparse
      # matrix. Note the dimensions of ind_log_sums_per_nest are (obs, nests).
      # Calculates sum _{j \in C_m} exp(V_{ij} / \lambda_m) for each nest m.
      ind_exp_sums_per_nest = (rows_to_obs.T *
                               rows_to_nests.multiply(exp_scaled_index[:, None]))
      # Ensure that ind_exp_sums_per_nest is an ndarray
      if isinstance(ind_exp_sums_per_nest, np.matrixlib.defmatrix.matrix):
          ind_exp_sums_per_nest = np.asarray(ind_exp_sums_per_nest)
      elif issparse(ind_exp_sums_per_nest):
          ind_exp_sums_per_nest = ind_exp_sums_per_nest.toarray()
      # Guard against overflow
      inf_idx = np.isposinf(ind_exp_sums_per_nest)
      ind_exp_sums_per_nest[inf_idx] = max_comp_value
  
      # Get the long-format representation of ind_log_sums_per_nest. Each row
      # will have two columns, one for each nest. The entries of the matrix will
      # be the log-sum for each nest, for the individual associated with the
      # given row. The "*" is used to perform the dot product since rows_to_obs
      # is a sparse matrix & ind_exp_sums_per_nest is a dense numpy matrix.
      long_exp_sums_per_nest = rows_to_obs.dot(ind_exp_sums_per_nest)
      if isinstance(long_exp_sums_per_nest, np.matrixlib.defmatrix.matrix):
          long_exp_sums_per_nest = np.asarray(long_exp_sums_per_nest)
  
      # Get the relevant log-sum for each row of the long-format data
      # Note the .A converts the numpy matrix into a numpy array
      # This is sum _{j \in C_m} exp(V_{ij} / \lambda_m) for the nest
      # belonging to each row
      long_exp_sums = (rows_to_nests.multiply(long_exp_sums_per_nest)
                                    .sum(axis=1)
                                    .A).ravel()
  
      # Get the denominators for each individual
      ind_denom = (np.power(ind_exp_sums_per_nest,
                            nest_coefs[None, :])
                     .sum(axis=1))
      # Guard against overflow and underflow
      inf_idx = np.isposinf(ind_denom)
      ind_denom[inf_idx] = max_comp_value
  
      zero_idx = (ind_denom == 0)
      ind_denom[zero_idx] = min_comp_value
  
      # Get the long format denominators.
      long_denom = rows_to_obs.dot(ind_denom)
      # Ensure that long_denom is 1D.
      long_denom.ravel()
  
      # Get the long format numerators
      long_numerators = (exp_scaled_index *
                         np.power(long_exp_sums,
                                  (long_nest_coefs - 1)))
      # Guard agains overflow and underflow
      inf_idx = np.isposinf(long_numerators)
      long_numerators[inf_idx] = max_comp_value
  
      zero_idx = (long_numerators == 0)
      long_numerators[zero_idx] = min_comp_value
  
      # Calculate and return the long-format probabilities
      long_probs = (long_numerators / long_denom).ravel()
      # Guard against underflow
      long_probs[np.where(long_probs == 0)] = min_comp_value
  
      # If desired, isolate the probabilities of the chosen alternatives
      if chosen_row_to_obs is None:
          chosen_probs = None
      else:
          # chosen_probs will be of shape (num_observations,)
          chosen_probs = (chosen_row_to_obs.transpose()
                                           .dot(long_probs))
          chosen_probs = np.asarray(chosen_probs).ravel()
  
      # Return the long form and chosen probabilities if desired
      if return_type == 'long_and_chosen_probs':
          return chosen_probs, long_probs
      # If working with predictions, return just the long form probabilities
      elif return_type == 'long_probs':
          return long_probs
      # If estimating the model and storing fitted probabilities or testing the
      # model on data for which we know the chosen alternative, just return the
      # chosen probabilities.
      elif return_type == 'chosen_probs':
          return chosen_probs
      # If we want all the factors of the probability (e.g. as when calculating
      # the gradient)
      elif return_type == 'all_prob_dict':
          # Create the dictionary of the various probabilities to be returned
          prob_dict = {}
          prob_dict["long_probs"] = long_probs
          prob_dict["chosen_probs"] = chosen_probs
  
          # Calculate the 'prob_given_nest' array
          prob_given_nest = exp_scaled_index / long_exp_sums
          # Guard against underflow
          zero_idx = (prob_given_nest == 0)
          prob_given_nest[zero_idx] = min_comp_value
  
          # Calculate the 'nest_choice_probs'. Note ind_denom is a matrix with
          # shape (num_obs, 1) so no need to explicitly broadcast
          nest_choice_probs = (np.power(ind_exp_sums_per_nest,
                                        nest_coefs[None, :]) /
                               ind_denom[:, None])
          # Guard against underflow
          zero_idx = (nest_choice_probs == 0)
          nest_choice_probs[zero_idx] = min_comp_value
          # Return dictionary.
          # Note the ".A" converts the numpy matrix into a numpy array
          prob_dict["prob_given_nest"] = prob_given_nest
          prob_dict["nest_choice_probs"] = nest_choice_probs
          prob_dict["ind_sums_per_nest"] = ind_exp_sums_per_nest
  
          return prob_dict