aspire.utils

aspire.utils#

Attributes#

`logger`
`IS_NAMESPACE_FUNCTIONS`

Classes#

`PoolHandler`	Context manager to temporarily replace the log_likelihood method of a
`AspireFile`	A subclass of h5py.File that adds metadata to the file.
`CallHistory`	Class to store the history of calls to a function.

Functions#

`configure_logger`([log_level, log_file, ...])	Configure the logger.
`logit`(x[, eps])	Logit function that also returns log Jacobian determinant.
`sigmoid`(x[, eps])	Sigmoid function that also returns log Jacobian determinant.
`logsumexp`(x[, axis])	Implementation of logsumexp that works with array api.
`to_numpy`(x, **kwargs)	Convert an array to a numpy array.
`asarray`(x[, xp, dtype])	Convert an array to the specified array API.
`determine_backend_name`([x, xp])	Determine the backend name from an array or array API module.
`resolve_dtype`(dtype, xp)	Resolve a dtype specification into an XP-specific dtype.
`convert_dtype`(dtype, target_xp, *[, source_xp])	Convert a dtype between array API namespaces.
`copy_array`(x[, xp])	Copy an array based on the array API being used.
`effective_sample_size`(log_w)
`disable_gradients`(xp[, inference])	Disable gradients for a specific array API.
`encode_dtype`(xp, dtype)	Encode a dtype for storage in an HDF5 file.
`decode_dtype`(xp, encoded_dtype)	Decode a dtype from an HDF5 file.
`encode_samples`(samples)	Encode a BaseSamples object for storage in an HDF5 file.
`decode_samples`(encoded_samples)	Decode a BaseSamples object from a dictionary loaded from an HDF5 file.
`encode_for_hdf5`(value)	Encode a value for storage in an HDF5 file.
`decode_from_hdf5`(value)	Decode a value loaded from an HDF5 file, reversing encode_for_hdf5.
`dump_pickle_to_hdf`(memfp, fp[, path, dsetname])	Dump pickled data to an HDF5 file object.
`dump_state`(state, fp[, path, dsetname, protocol])	Pickle a state object and store it in an HDF5 dataset.
`resolve_xp`(xp_name)	Resolve a backend name to the corresponding array_api_compat module.
`infer_device`(x, xp)	Best-effort device inference that avoids non-portable identifiers.
`safe_to_device`(x, device, xp)	Move to device if specified; otherwise return input.
`recursively_save_to_h5_file`(h5_file, path, dictionary)	Save a dictionary to an HDF5 file with flattened keys under a given group path.
`load_from_h5_file`(h5_file, path)	Load a flattened dictionary from an HDF5 group and rebuild nesting.
`get_package_version`(package_name)	Get the version of a package.
`update_at_indices`(x, slc, y)	Update an array at specific indices."
`track_calls`([wrapped])	Decorator to track calls to a function.
`function_id`(fn)	Get a unique identifier for a function.
`enable_scipy_array_api`()	Set the environment variable to enable array API compatibility in SciPy.

Module Contents#

aspire.utils.logger[source]#

aspire.utils.IS_NAMESPACE_FUNCTIONS[source]#

aspire.utils.configure_logger(log_level='INFO', log_file=None, additional_loggers=None, include_aspire_loggers=True, stream=None)[source]#

Configure the logger.

Adds a stream handler to the logger and optionally a file handler.

Parameters:

log_level (str or int, optional) – The log level to use. Defaults to “INFO”.
log_file (str or Path, optional) – The file to write logs to. If None, logs are not written to a file.
additional_loggers (list of str, optional) – Additional loggers to configure. Defaults to None.
include_aspire_loggers (bool, optional) – Whether to include all loggers that start with “aspire_” or “aspire-“. Defaults to True.
stream (str, optional) – The stream to use for the stream handler. Defaults to None, which uses sys.stderr.

Returns:

The configured logger.

Return type:

logging.Logger

class aspire.utils.PoolHandler(aspire_instance, pool, close_pool=True, parallelize_prior=False)[source]#

Context manager to temporarily replace the log_likelihood method of a aspire instance with a version that uses a multiprocessing pool to parallelize computation.

Parameters:

aspire_instance (aspire) – The aspire instance to modify. The log_likelihood method of this instance must accept a map_fn keyword argument.
pool (multiprocessing.Pool) – The pool to use for parallel computation.
close_pool (bool, optional) – Whether to close the pool when exiting the context manager. Defaults to True.
parallelize_prior (bool, optional) – 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 map_fn keyword argument.

parallelize_prior = False[source]#

property aspire_instance[source]#

pool[source]#

close_pool = True[source]#

aspire.utils.logit(x, eps=None)[source]#

Logit function that also returns log Jacobian determinant.

Parameters:

x (float or ndarray) – Array of values
eps (float, optional) – Epsilon value used to clamp inputs to [eps, 1 - eps]. If None, then inputs are not clamped.

Returns:

float or ndarray – Rescaled values.
float or ndarray – Log Jacobian determinant.

Return type:

tuple[array_api_compat.common._typing.Array, array_api_compat.common._typing.Array]

aspire.utils.sigmoid(x, eps=None)[source]#

Sigmoid function that also returns log Jacobian determinant.

Parameters:

x (float or ndarray) – Array of values
eps (float, optional) – Epsilon value used to clamp inputs to [eps, 1 - eps]. If None, then inputs are not clamped.

Returns:

float or ndarray – Rescaled values.
float or ndarray – Log Jacobian determinant.

Return type:

tuple[array_api_compat.common._typing.Array, array_api_compat.common._typing.Array]

aspire.utils.logsumexp(x, axis=None)[source]#

Implementation of logsumexp that works with array api.

This will be removed once the implementation in scipy is compatible.

Parameters:

x (array_api_compat.common._typing.Array)
axis (int | None)

Return type:

array_api_compat.common._typing.Array

aspire.utils.to_numpy(x, **kwargs)[source]#

Convert an array to a numpy array.

This automatically moves the array to the CPU.

Parameters:

x (Array) – The array to convert.
kwargs (dict) – Additional keyword arguments to pass to numpy.asarray.

Return type:

array_api_compat.numpy.ndarray

aspire.utils.asarray(x, xp=None, dtype=None, **kwargs)[source]#

Convert an array to the specified array API.

Parameters:

x (Array) – The array to convert.
xp (Any) – The array API to use for the conversion. If None, the array API is inferred from the input array.
dtype (Any | str | None) – The dtype to use for the conversion. If None, the dtype is not changed.
kwargs (dict) – Additional keyword arguments to pass to xp.asarray.

Return type:

array_api_compat.common._typing.Array

aspire.utils.determine_backend_name(x=None, xp=None)[source]#

Determine the backend name from an array or array API module.

Parameters:

x (Array or None) – The array to infer the backend from. If None, xp must be provided.
xp (Any or None) – The array API module to infer the backend from. If None, x must be provided.

Returns:

The name of the backend. If the backend cannot be determined, returns “unknown”.

Return type:

str

aspire.utils.resolve_dtype(dtype, xp)[source]#

Resolve a dtype specification into an XP-specific dtype.

Parameters:

dtype (Any | str | None) – The dtype specification. Can be None, a string, or a dtype-like object.
xp (module) – The array API module that should interpret the dtype.

Returns:

The resolved dtype object compatible with xp (or None if unspecified).

Return type:

Any | None

aspire.utils.convert_dtype(dtype, target_xp, *, source_xp=None)[source]#

Convert a dtype between array API namespaces.

Parameters:

dtype (Any | str | None) – The dtype to convert. Can be a dtype object, string, or None.
target_xp (module) – The target array API namespace to convert the dtype into.
source_xp (module, optional) – The source namespace of the dtype. Provided for API symmetry and future use; currently unused but accepted.

Returns:

The dtype object compatible with target_xp (or None if dtype is None).

Return type:

Any | None

aspire.utils.copy_array(x, xp=None)[source]#

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.

Parameters:

x (Array) – The array to copy.
xp (Any) – The array API to use for the copy.

Returns:

The copied array.

Return type:

Array

aspire.utils.effective_sample_size(log_w)[source]#

Parameters:: log_w (array_api_compat.common._typing.Array)
Return type:: float

aspire.utils.disable_gradients(xp, inference=True)[source]#

Disable gradients for a specific array API.

Usage:

```python with disable_gradients(xp):

# Do something

```

Parameters:

xp (module) – The array API module to use.
inference (bool, optional) – When using PyTorch, set to True to enable inference mode.

aspire.utils.encode_dtype(xp, dtype)[source]#

Encode a dtype for storage in an HDF5 file.

Parameters:

xp (module) – The array API module to use.
dtype (dtype) – The dtype to encode.

Returns:

The encoded dtype.

Return type:

str

aspire.utils.decode_dtype(xp, encoded_dtype)[source]#

Decode a dtype from an HDF5 file.

Parameters:

xp (module) – The array API module to use.
encoded_dtype (dict) – The encoded dtype.

Returns:

The decoded dtype.

Return type:

dtype

aspire.utils.encode_samples(samples)[source]#

Encode a BaseSamples object for storage in an HDF5 file.

Parameters:: samples (BaseSamples) – The samples to encode.
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.
Return type:: dict

aspire.utils.decode_samples(encoded_samples)[source]#

Decode a BaseSamples object from a dictionary loaded from an HDF5 file.

Parameters:: encoded_samples (dict) – The dictionary containing the encoded samples. This should have been produced by the encode_samples function.
Raises:: ValueError – If the encoded_samples dictionary does not have the expected format or if the samples type is unknown.
Return type:: aspire.samples.BaseSamples

aspire.utils.encode_for_hdf5(value)[source]#

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 encode_samples

Parameters:: value (Any)
Return type:: Any

aspire.utils.decode_from_hdf5(value)[source]#

Decode a value loaded from an HDF5 file, reversing encode_for_hdf5.

Parameters:: value (Any)
Return type:: Any

aspire.utils.dump_pickle_to_hdf(memfp, fp, path=None, dsetname='state')[source]#

Dump pickled data to an HDF5 file object.

Parameters:

memfp (BytesIO) – A BytesIO object containing the pickled data.
fp (h5py.File) – An open h5py.File object to write to.
path (str, optional) – The group path within the HDF5 file to store the dataset. If None, the dataset is stored at the root level. Defaults to None.
dsetname (str, optional) – The name of the dataset to create or overwrite. Defaults to “state”.

aspire.utils.dump_state(state, fp, path=None, dsetname='state', protocol=pickle.HIGHEST_PROTOCOL)[source]#: Pickle a state object and store it in an HDF5 dataset.

aspire.utils.resolve_xp(xp_name)[source]#

Resolve a backend name to the corresponding array_api_compat module.

Returns None if the name is None or cannot be resolved.

Parameters:: xp_name (str | None)

aspire.utils.infer_device(x, xp)[source]#

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.

aspire.utils.safe_to_device(x, device, xp)[source]#

Move to device if specified; otherwise return input.

Skips moves for numpy/jax/None devices; logs and returns input on failure.

aspire.utils.recursively_save_to_h5_file(h5_file, path, dictionary)[source]#: Save a dictionary to an HDF5 file with flattened keys under a given group path.

aspire.utils.load_from_h5_file(h5_file, path)[source]#: Load a flattened dictionary from an HDF5 group and rebuild nesting.

aspire.utils.get_package_version(package_name)[source]#

Get the version of a package.

Parameters:: package_name (str) – The name of the package.
Returns:: The version of the package.
Return type:: str

class aspire.utils.AspireFile(*args, **kwargs)[source]#

Bases: h5py.File

A subclass of h5py.File that adds metadata to the file.

aspire.utils.update_at_indices(x, slc, y)[source]#

Update an array at specific indices.”

This is a workaround for the fact that array API does not support advanced indexing with all backends.

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]]

Parameters:

x (Array) – The array to update.
slc (Array) – The indices to update.
y (Array) – The values to set at the indices.

Returns:

The updated array.

Return type:

Array

class aspire.utils.CallHistory[source]#

Class to store the history of calls to a function.

args[source]#

The positional arguments of each call.

Type:: list[tuple]

kwargs[source]#

The keyword arguments of each call.

Type:: list[dict]

args: list[tuple][source]#

kwargs: list[dict][source]#

to_dict(list_to_dict=False)[source]#

Convert the call history to a dictionary.

Parameters:: list_to_dict (bool) – 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.
Return type:: dict[str, Any]

aspire.utils.track_calls(wrapped=None)[source]#

Decorator to track calls to a function.

The decorator adds a calls attribute to the wrapped function, which is a CallHistory object that stores the arguments and keyword arguments of each call.

aspire.utils.function_id(fn)[source]#

Get a unique identifier for a function.

Parameters:: fn (Any) – The function to get the identifier for.
Returns:: The unique identifier for the function.
Return type:: str

aspire.utils.enable_scipy_array_api()[source]#

Set the environment variable to enable array API compatibility in SciPy.

Set SCIPY_ARRAY_API=1 if it is not already set.

Return type:: None