docs: def calc_nested_gradient(orig_nest_coefs,
                           index_coefs,
                           design,
                           choice_vec,
                           rows_to_obs,
                           rows_to_nests,
                           ridge=None,
                           weights=None,
                           use_jacobian=True,
                           *args,
                           **kwargs):
      """
      Parameters
      ----------
      orig_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 logit of
          the inverse of the scale coefficients for each 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 choosing each
          alternative. 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.
      choice_vec : 1D ndarray.
          All elements should by ints, floats, or longs. Each element represents
          whether the individual associated with the given row chose the
          alternative associated with the given row.
      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 be 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).
      ridge : int, float, long, or None, optional.
          Determines whether or not ridge regression is performed. If an int,
          float or long is passed, then that scalar determines the ridge penalty
          for the optimization. Default `== None`.
      weights : 1D ndarray or None.
          Allows for the calculation of weighted log-likelihoods. The weights can
          represent various things. In stratified samples, the weights may be
          the proportion of the observations in a given strata for a sample in
          relation to the proportion of observations in that strata in the
          population. In latent class models, the weights may be the probability
          of being a particular class.
      use_jacobian : bool, optional.
          Determines whether or not the jacobian will be used when calculating
          the gradient. When performing model estimation, `use_jacobian` should
          be `True` if the values being estimated are actually the logit of the
          nest coefficients. Default `== True`.
  
      Returns
      -------
      gradient : 1D numpy array.
         The gradient of the log-likelihood with respect to the given nest
         coefficients and index coefficients.
      """
      # Calculate the weights for the sample
      if weights is None:
          weights = np.ones(design.shape[0])
      weights_per_obs = np.max(rows_to_obs.toarray() * weights[:, None], axis=0)
  
      # Transform the nest coefficients into their "always positive" versions
      nest_coefs = naturalize_nest_coefs(orig_nest_coefs)
  
      # Get the vectors and matrices needed to calculate the gradient
      vector_dict = prep_vectors_for_gradient(nest_coefs,
                                              index_coefs,
                                              design,
                                              choice_vec,
                                              rows_to_obs,
                                              rows_to_nests)
  
      # Calculate the index for each alternative for each person
      sys_utility = design.dot(index_coefs)
  
      # Calculate w_ij
      long_w = sys_utility / vector_dict["long_nest_params"]
      # Guard against overflow
      inf_index = np.isposinf(long_w)
      long_w[inf_index] = max_comp_value
  
      ##########
      # Calculate d_log_likelihood_d_nest_params
      ##########
      # Calculate the term that onlny depends on nest level values
      log_exp_sums = np.log(vector_dict["ind_sums_per_nest"])
      # Guard against overflow
      log_exp_sums[np.isneginf(log_exp_sums)] = -1 * max_comp_value
  
      # Calculate the first term of the derivative of the log-liikelihood
      # with respect to the nest parameters
      nest_gradient_term_1 = ((vector_dict["obs_to_chosen_nests"] -
                               vector_dict["nest_choice_probs"]) *
                              log_exp_sums *
                              weights_per_obs[:, None]).sum(axis=0)
  
      # Calculate the second term of the derivative of the log-liikelihood
      # with respect to the nest parameters
      half_deriv = ((vector_dict["long_probs"] -
                     vector_dict["long_chosen_nest"] *
                     vector_dict["prob_given_nest"]) *
                    long_w *
                    weights)
      nest_gradient_term_2 = (rows_to_nests.transpose()
                                           .dot(half_deriv)[:, None]).ravel()
  
      # Calculate the third term of the derivative of the log-likelihood
      # with respect to the nest parameters
      nest_gradient_term_3a = (choice_vec -
                               vector_dict["long_chosen_nest"] *
                               vector_dict["prob_given_nest"])
      nest_gradient_term_3b = ((-1 * nest_gradient_term_3a * long_w * weights) /
                               vector_dict["long_nest_params"])
      # Guard against overflow
      inf_idx = np.isposinf(nest_gradient_term_3b)
      nest_gradient_term_3b[inf_idx] = max_comp_value
  
      neg_inf_idx = np.isneginf(nest_gradient_term_3b)
      nest_gradient_term_3b[neg_inf_idx] = -1 * max_comp_value
  
      # Get the nest-wide version of this piece of the gradient
      nest_gradient_term_3 = (rows_to_nests.transpose()
                                           .dot(nest_gradient_term_3b)).ravel()
  
      # Combine the two terms. Note the "nest_coefs * (1 - nest_coefs)" is due to
      # the fact that we're estimating the logit of the nest coefficients instead
      # of the nest coefficient itself. We therefore need to multiply by
      # d_nest_coef_d_estimated_variable to get the correct gradient.
      # d_nest_coef_d_estimated_variable == nest_coefs * (1 - nest_coefs).
      if use_jacobian:
          jacobian = nest_coefs * (1.0 - nest_coefs)
      else:
          jacobian = 1
      nest_gradient = ((nest_gradient_term_1 +
                        nest_gradient_term_2 +
                        nest_gradient_term_3) *
                       jacobian)[None, :]
  
      ##########
      # Calculate d_loglikelihood_d_beta
      ##########
      beta_gradient_term_1 = ((vector_dict["scaled_y"] -
                               vector_dict["p_tilde_given_nest"] +
                               vector_dict["p_tilde_given_nest"] *
                               vector_dict["long_nest_params"] -
                               vector_dict["long_probs"]) *
                              weights)[None, :]
      #####
      # Calculate the derivative with respect to beta
      #####
      beta_gradient = beta_gradient_term_1.dot(design)
  
      #####
      # Combine the gradient pieces and account for ridge parameter
      #####
      gradient = np.concatenate((nest_gradient, beta_gradient), axis=1).ravel()
  
      if ridge is not None:
          # Note that the 20 is used in place of 'infinity' since I would really
          # like to specify the expected value of the nest coefficient to 1, but
          # that would make the logit of the nest parameter infinity. Instead I
          # use 20 as a close enough value-- (1 + exp(-20))**-1 is approx. 1.
          params = np.concatenate(((20 - orig_nest_coefs), index_coefs), axis=0)
  
          gradient -= 2 * ridge * params
  
      return gradient