original_docstring: Calculate the array of sigma_points for the unscented transform.
  
      Args:
          states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
              pre-update states estimates.
          upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
              n_states) with the transpose of the lower triangular cholesky factor
              of the pre-update covariance matrix of the state estimates.
          scaling_factor (float): A scaling factor that controls the spread of the
              sigma points. Bigger means that sigma points are further apart. Depends on
              the sigma_point algorithm chosen.
  
      Returns:
          jax.numpy.array: Array of shape n_obs, n_mixtures, n_sigma, n_fac (where n_sigma
          equals 2 * n_fac + 1) with sigma points.
  
      
docstring: Calculate the array of sigma_points for the unscented transform.
docstring_tokens: Calculate the array of sigma_points for the unscented transform .
code: def _calculate_sigma_points(states, upper_chols, scaling_factor):
      """Calculate the array of sigma_points for the unscented transform.
  
      Args:
          states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
              pre-update states estimates.
          upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
              n_states) with the transpose of the lower triangular cholesky factor
              of the pre-update covariance matrix of the state estimates.
          scaling_factor (float): A scaling factor that controls the spread of the
              sigma points. Bigger means that sigma points are further apart. Depends on
              the sigma_point algorithm chosen.
  
      Returns:
          jax.numpy.array: Array of shape n_obs, n_mixtures, n_sigma, n_fac (where n_sigma
          equals 2 * n_fac + 1) with sigma points.
  
      """
      n_obs, n_mixtures, n_fac = states.shape
      n_sigma = 2 * n_fac + 1
  
      scaled_upper_chols = upper_chols * scaling_factor
      sigma_points = jnp.repeat(states, n_sigma, axis=1).reshape(
          n_obs, n_mixtures, n_sigma, n_fac
      )
      sigma_points = index_add(
          sigma_points, index[:, :, 1 : n_fac + 1], scaled_upper_chols
      )
      sigma_points = index_add(
          sigma_points, index[:, :, n_fac + 1 :], -scaled_upper_chols
      )
      return sigma_points
code_tokens: def _calculate_sigma_points ( states , upper_chols , scaling_factor ) : n_obs , n_mixtures , n_fac = states . shape n_sigma = 2 * n_fac + 1 scaled_upper_chols = upper_chols * scaling_factor sigma_points = jnp . repeat ( states , n_sigma , axis = 1 ) . reshape ( n_obs , n_mixtures , n_sigma , n_fac ) sigma_points = index_add ( sigma_points , index [ : , : , 1 : n_fac + 1 ] , scaled_upper_chols ) sigma_points = index_add ( sigma_points , index [ : , : , n_fac + 1 : ] , - scaled_upper_chols ) return sigma_points
short_docstring: Calculate the array of sigma_points for the unscented transform.
short_docstring_tokens: Calculate the array of sigma_points for the unscented transform .
comment: """Calculate the array of sigma_points for the unscented transform.
  
      Args:
          states (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states) with
              pre-update states estimates.
          upper_chols (jax.numpy.array): Array of shape (n_obs, n_mixtures, n_states,
              n_states) with the transpose of the lower triangular cholesky factor
              of the pre-update covariance matrix of the state estimates.
          scaling_factor (float): A scaling factor that controls the spread of the
              sigma points. Bigger means that sigma points are further apart. Depends on
              the sigma_point algorithm chosen.
  
      Returns:
          jax.numpy.array: Array of shape n_obs, n_mixtures, n_sigma, n_fac (where n_sigma
          equals 2 * n_fac + 1) with sigma points.
  
      """