aspire.utils ============ .. py:module:: aspire.utils Attributes ---------- .. autoapisummary:: aspire.utils.logger aspire.utils.IS_NAMESPACE_FUNCTIONS Classes ------- .. autoapisummary:: aspire.utils.PoolHandler aspire.utils.AspireFile aspire.utils.CallHistory Functions --------- .. autoapisummary:: aspire.utils.configure_logger aspire.utils.logit aspire.utils.sigmoid aspire.utils.logsumexp aspire.utils.to_numpy aspire.utils.asarray aspire.utils.determine_backend_name aspire.utils.resolve_dtype aspire.utils.convert_dtype aspire.utils.copy_array aspire.utils.effective_sample_size aspire.utils.disable_gradients aspire.utils.encode_dtype aspire.utils.decode_dtype aspire.utils.encode_samples aspire.utils.decode_samples aspire.utils.encode_for_hdf5 aspire.utils.decode_from_hdf5 aspire.utils.dump_pickle_to_hdf aspire.utils.dump_state aspire.utils.resolve_xp aspire.utils.infer_device aspire.utils.safe_to_device aspire.utils.recursively_save_to_h5_file aspire.utils.load_from_h5_file aspire.utils.get_package_version aspire.utils.update_at_indices aspire.utils.track_calls aspire.utils.function_id aspire.utils.enable_scipy_array_api Module Contents --------------- .. py:data:: logger .. py:data:: IS_NAMESPACE_FUNCTIONS .. py:function:: configure_logger(log_level = 'INFO', log_file = None, additional_loggers = None, include_aspire_loggers = True, stream = None) Configure the logger. Adds a stream handler to the logger and optionally a file handler. :param log_level: The log level to use. Defaults to "INFO". :type log_level: :py:class:`str` or :py:class:`int`, *optional* :param log_file: The file to write logs to. If None, logs are not written to a file. :type log_file: :py:class:`str` or :py:class:`Path`, *optional* :param additional_loggers: Additional loggers to configure. Defaults to None. :type additional_loggers: :py:class:`list` of :py:class:`str`, *optional* :param include_aspire_loggers: Whether to include all loggers that start with "aspire_" or "aspire-". Defaults to True. :type include_aspire_loggers: :py:class:`bool`, *optional* :param stream: The stream to use for the stream handler. Defaults to None, which uses sys.stderr. :type stream: :py:class:`str`, *optional* :returns: The configured logger. :rtype: :py:class:`logging.Logger` .. py:class:: PoolHandler(aspire_instance, pool, close_pool = True, parallelize_prior = False) Context manager to temporarily replace the log_likelihood method of a aspire instance with a version that uses a multiprocessing pool to parallelize computation. :param aspire_instance: The aspire instance to modify. The log_likelihood method of this instance must accept a :code:`map_fn` keyword argument. :type aspire_instance: :py:class:`aspire` :param pool: The pool to use for parallel computation. :type pool: :py:class:`multiprocessing.Pool` :param close_pool: Whether to close the pool when exiting the context manager. Defaults to True. :type close_pool: :py:class:`bool`, *optional* :param parallelize_prior: Whether to parallelize the log_prior method as well. Defaults to False. If True, the log_prior method of the aspire instance must also accept a :code:`map_fn` keyword argument. :type parallelize_prior: :py:class:`bool`, *optional* .. py:attribute:: parallelize_prior :value: False .. py:property:: aspire_instance .. py:attribute:: pool .. py:attribute:: close_pool :value: True .. py:function:: logit(x, eps = None) Logit function that also returns log Jacobian determinant. :param x: Array of values :type x: :py:class:`float` or :py:class:`ndarray` :param eps: Epsilon value used to clamp inputs to [eps, 1 - eps]. If None, then inputs are not clamped. :type eps: :py:class:`float`, *optional* :returns: * :py:class:`float` or :py:class:`ndarray` -- Rescaled values. * :py:class:`float` or :py:class:`ndarray` -- Log Jacobian determinant. .. py:function:: sigmoid(x, eps = None) Sigmoid function that also returns log Jacobian determinant. :param x: Array of values :type x: :py:class:`float` or :py:class:`ndarray` :param eps: Epsilon value used to clamp inputs to [eps, 1 - eps]. If None, then inputs are not clamped. :type eps: :py:class:`float`, *optional* :returns: * :py:class:`float` or :py:class:`ndarray` -- Rescaled values. * :py:class:`float` or :py:class:`ndarray` -- Log Jacobian determinant. .. py:function:: logsumexp(x, axis = None) Implementation of logsumexp that works with array api. This will be removed once the implementation in scipy is compatible. .. py:function:: to_numpy(x, **kwargs) Convert an array to a numpy array. This automatically moves the array to the CPU. :param x: The array to convert. :type x: :py:class:`Array` :param kwargs: Additional keyword arguments to pass to numpy.asarray. :type kwargs: :py:class:`dict` .. py:function:: asarray(x, xp = None, dtype = None, **kwargs) Convert an array to the specified array API. :param x: The array to convert. :type x: :py:class:`Array` :param xp: The array API to use for the conversion. If None, the array API is inferred from the input array. :type xp: :py:class:`Any` :param dtype: The dtype to use for the conversion. If None, the dtype is not changed. :type dtype: :py:class:`Any | str | None` :param kwargs: Additional keyword arguments to pass to xp.asarray. :type kwargs: :py:class:`dict` .. py:function:: determine_backend_name(x = None, xp = None) Determine the backend name from an array or array API module. :param x: The array to infer the backend from. If None, xp must be provided. :type x: :py:class:`Array` or :py:obj:`None` :param xp: The array API module to infer the backend from. If None, x must be provided. :type xp: :py:class:`Any` or :py:obj:`None` :returns: The name of the backend. If the backend cannot be determined, returns "unknown". :rtype: :py:class:`str` .. py:function:: resolve_dtype(dtype, xp) Resolve a dtype specification into an XP-specific dtype. :param dtype: The dtype specification. Can be None, a string, or a dtype-like object. :type dtype: :py:class:`Any | str | None` :param xp: The array API module that should interpret the dtype. :type xp: :py:class:`module` :returns: The resolved dtype object compatible with ``xp`` (or None if unspecified). :rtype: :py:class:`Any | None` .. py:function:: convert_dtype(dtype, target_xp, *, source_xp = None) Convert a dtype between array API namespaces. :param dtype: The dtype to convert. Can be a dtype object, string, or None. :type dtype: :py:class:`Any | str | None` :param target_xp: The target array API namespace to convert the dtype into. :type target_xp: :py:class:`module` :param source_xp: The source namespace of the dtype. Provided for API symmetry and future use; currently unused but accepted. :type source_xp: :py:class:`module`, *optional* :returns: The dtype object compatible with ``target_xp`` (or None if ``dtype`` is None). :rtype: :py:class:`Any | None` .. py:function:: copy_array(x, xp = None) Copy an array based on the array API being used. This uses the most appropriate method to copy the array depending on the array API. :param x: The array to copy. :type x: :py:class:`Array` :param xp: The array API to use for the copy. :type xp: :py:class:`Any` :returns: The copied array. :rtype: :py:class:`Array` .. py:function:: effective_sample_size(log_w) .. py:function:: disable_gradients(xp, inference = True) Disable gradients for a specific array API. Usage: ```python with disable_gradients(xp): # Do something ``` :param xp: The array API module to use. :type xp: :py:class:`module` :param inference: When using PyTorch, set to True to enable inference mode. :type inference: :py:class:`bool`, *optional* .. py:function:: encode_dtype(xp, dtype) Encode a dtype for storage in an HDF5 file. :param xp: The array API module to use. :type xp: :py:class:`module` :param dtype: The dtype to encode. :type dtype: :py:class:`dtype` :returns: The encoded dtype. :rtype: :py:class:`str` .. py:function:: decode_dtype(xp, encoded_dtype) Decode a dtype from an HDF5 file. :param xp: The array API module to use. :type xp: :py:class:`module` :param encoded_dtype: The encoded dtype. :type encoded_dtype: :py:class:`dict` :returns: The decoded dtype. :rtype: :py:class:`dtype` .. py:function:: encode_samples(samples) Encode a BaseSamples object for storage in an HDF5 file. :param samples: The samples to encode. :type samples: :py:class:`BaseSamples` :returns: A dictionary containing the encoded samples. This includes a marker to indicate that it is an encoded samples object, the type of the samples, and the data needed to reconstruct the samples. :rtype: :py:class:`dict` .. py:function:: decode_samples(encoded_samples) Decode a BaseSamples object from a dictionary loaded from an HDF5 file. :param encoded_samples: The dictionary containing the encoded samples. This should have been produced by the :code:`encode_samples` function. :type encoded_samples: :py:class:`dict` :raises ValueError: If the encoded_samples dictionary does not have the expected format or if the samples type is unknown. .. py:function:: encode_for_hdf5(value) Encode a value for storage in an HDF5 file. Special cases: - None is replaced with "__none__" - Empty dictionaries are replaced with "__empty_dict__" - BaseSamples objects are encoded using :code:`encode_samples` .. py:function:: decode_from_hdf5(value) Decode a value loaded from an HDF5 file, reversing encode_for_hdf5. .. py:function:: dump_pickle_to_hdf(memfp, fp, path=None, dsetname='state') Dump pickled data to an HDF5 file object. :param memfp: A BytesIO object containing the pickled data. :type memfp: :py:class:`BytesIO` :param fp: An open h5py.File object to write to. :type fp: :py:class:`h5py.File` :param path: The group path within the HDF5 file to store the dataset. If None, the dataset is stored at the root level. Defaults to None. :type path: :py:class:`str`, *optional* :param dsetname: The name of the dataset to create or overwrite. Defaults to "state". :type dsetname: :py:class:`str`, *optional* .. py:function:: dump_state(state, fp, path=None, dsetname='state', protocol=pickle.HIGHEST_PROTOCOL) Pickle a state object and store it in an HDF5 dataset. .. py:function:: resolve_xp(xp_name) Resolve a backend name to the corresponding array_api_compat module. Returns None if the name is None or cannot be resolved. .. py:function:: infer_device(x, xp) Best-effort device inference that avoids non-portable identifiers. Returns None for numpy/jax backends; returns the backend device object for torch/cupy if available. .. py:function:: safe_to_device(x, device, xp) Move to device if specified; otherwise return input. Skips moves for numpy/jax/None devices; logs and returns input on failure. .. py:function:: recursively_save_to_h5_file(h5_file, path, dictionary) Save a dictionary to an HDF5 file with flattened keys under a given group path. .. py:function:: load_from_h5_file(h5_file, path) Load a flattened dictionary from an HDF5 group and rebuild nesting. .. py:function:: get_package_version(package_name) Get the version of a package. :param package_name: The name of the package. :type package_name: :py:class:`str` :returns: The version of the package. :rtype: :py:class:`str` .. py:class:: AspireFile(*args, **kwargs) Bases: :py:obj:`h5py.File` A subclass of h5py.File that adds metadata to the file. .. py:function:: update_at_indices(x, slc, y) Update an array at specific indices." This is a workaround for the fact that array API does not support advanced indexing with all backends. .. rubric:: Examples >>> x = xp.array([[1, 2], [3, 4], [5, 6]]) >>> update_at_indices(x, (slice(None), 0), xp.array([10, 20, 30])) [[10 2] [20 4] [30 6]] :param x: The array to update. :type x: :py:class:`Array` :param slc: The indices to update. :type slc: :py:class:`Array` :param y: The values to set at the indices. :type y: :py:class:`Array` :returns: The updated array. :rtype: :py:class:`Array` .. py:class:: CallHistory Class to store the history of calls to a function. .. attribute:: args The positional arguments of each call. :type: :py:class:`list[tuple]` .. attribute:: kwargs The keyword arguments of each call. :type: :py:class:`list[dict]` .. py:attribute:: args :type: list[tuple] .. py:attribute:: kwargs :type: list[dict] .. py:method:: to_dict(list_to_dict = False) Convert the call history to a dictionary. :param list_to_dict: If True, convert the lists of args and kwargs to dictionaries with string keys. If False, keep them as lists. This is useful when encoding the history for HDF5. :type list_to_dict: :py:class:`bool` .. py:function:: track_calls(wrapped=None) Decorator to track calls to a function. The decorator adds a :code:`calls` attribute to the wrapped function, which is a :py:class:`CallHistory` object that stores the arguments and keyword arguments of each call. .. py:function:: function_id(fn) Get a unique identifier for a function. :param fn: The function to get the identifier for. :type fn: :py:class:`Any` :returns: The unique identifier for the function. :rtype: :py:class:`str` .. py:function:: enable_scipy_array_api() Set the environment variable to enable array API compatibility in SciPy. Set :code:`SCIPY_ARRAY_API=1` if it is not already set.