docs: def prep_vectors_for_gradient(nest_coefs,
                                index_coefs,
                                design,
                                choice_vec,
                                rows_to_obs,
                                rows_to_nests,
                                *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. Note, this
          is NOT THE LOGIT of the inverse of the scale coefficients.
      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. Should have the same number
          of rows as `design`.
      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).
  
      Returns
      -------
      desired_arrays : dict.
          Will contain the arrays necessary for calculating the gradient of the
          nested logit log-likelihood. The keys will be:
              `["long_nest_params", "scaled_y", "long_chosen_nest",
                "obs_to_chosen_nests", "p_tilde_given_nest", "long_probs",
                "prob_given_nest", "nest_choice_probs", "ind_sums_per_nest"]`
      """
      # Create the "long_nest_parameters" which is an array with one element per
      # alternative per observation, where each element is the nest parameter for
      # the alternative corresponding to the given row
      long_nest_params = (rows_to_nests.multiply(nest_coefs[None, :])
                                       .sum(axis=1)
                                       .A
                                       .ravel())
  
      # Calculate y-tilde
      scaled_y = choice_vec / long_nest_params
      # Guard against overflow
      inf_index = np.isinf(scaled_y)
      scaled_y[inf_index] = max_comp_value
  
      # Determine which nest was chosen by each row's individual.
      # Resulting matrix has shape (num_rows, num_nests)
      obs_to_chosen_nests = (rows_to_obs.T *
                             rows_to_nests.multiply(choice_vec[:, None])).A
      row_to_chosen_nest = rows_to_obs * obs_to_chosen_nests
      # Determine whether the given row is part of the nest that was chosen
      long_chosen_nest = (rows_to_nests.multiply(row_to_chosen_nest)
                                       .sum(axis=1)
                                       .A
                                       .ravel())
  
      # Get the various probabilities
      prob_dict = calc_nested_probs(nest_coefs,
                                    index_coefs,
                                    design,
                                    rows_to_obs,
                                    rows_to_nests,
                                    return_type='all_prob_dict')
  
      # Calculate p_tilde_ij_given_nest
      p_tilde_row_given_nest = (prob_dict["prob_given_nest"] *
                                long_chosen_nest /
                                long_nest_params)
      # Guard against overflow
      inf_index = np.isinf(p_tilde_row_given_nest)
      p_tilde_row_given_nest[inf_index] = max_comp_value
  
      # Return all the desired matrices and arrays
      desired_arrays = {}
      desired_arrays["long_nest_params"] = long_nest_params.ravel()
      desired_arrays["scaled_y"] = scaled_y.ravel()
      desired_arrays["long_chosen_nest"] = long_chosen_nest
      desired_arrays["obs_to_chosen_nests"] = obs_to_chosen_nests
      desired_arrays["p_tilde_given_nest"] = p_tilde_row_given_nest
      desired_arrays["long_probs"] = prob_dict["long_probs"]
      desired_arrays["prob_given_nest"] = prob_dict["prob_given_nest"]
      desired_arrays["nest_choice_probs"] = prob_dict["nest_choice_probs"]
      desired_arrays["ind_sums_per_nest"] = prob_dict["ind_sums_per_nest"]
  
      return desired_arrays