from __future__ import annotations
import functools
import importlib
import inspect
import logging
import pickle
import warnings
from contextlib import contextmanager
from dataclasses import dataclass
from functools import partial
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, Any
import array_api_compat.numpy as np
import array_api_extra as xpx
import h5py
import wrapt
from array_api_compat import (
array_namespace,
is_cupy_namespace,
is_dask_namespace,
is_jax_array,
is_jax_namespace,
is_ndonnx_namespace,
is_numpy_namespace,
is_pydata_sparse_namespace,
is_torch_array,
is_torch_namespace,
to_device,
)
if TYPE_CHECKING:
from multiprocessing import Pool
from array_api_compat.common._typing import Array
from .aspire import Aspire
from .samples import BaseSamples
[docs]
logger = logging.getLogger(__name__)
[docs]
IS_NAMESPACE_FUNCTIONS = {
"numpy": is_numpy_namespace,
"torch": is_torch_namespace,
"jax": is_jax_namespace,
"cupy": is_cupy_namespace,
"dask": is_dask_namespace,
"pydata_sparse": is_pydata_sparse_namespace,
"ndonnx": is_ndonnx_namespace,
}
[docs]
class PoolHandler:
"""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 :code:`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 :code:`map_fn` keyword argument.
"""
def __init__(
self,
aspire_instance: Aspire,
pool: Pool,
close_pool: bool = True,
parallelize_prior: bool = False,
):
[docs]
self.parallelize_prior = parallelize_prior
self.aspire_instance = aspire_instance
[docs]
self.close_pool = close_pool
@property
[docs]
def aspire_instance(self):
return self._aspire_instance
@aspire_instance.setter
def aspire_instance(self, value: Aspire):
signature = inspect.signature(value.log_likelihood)
if "map_fn" not in signature.parameters:
raise ValueError(
"The log_likelihood method of the Aspire instance must accept a"
" 'map_fn' keyword argument."
)
signature = inspect.signature(value.log_prior)
if "map_fn" not in signature.parameters and self.parallelize_prior:
raise ValueError(
"The log_prior method of the Aspire instance must accept a"
" 'map_fn' keyword argument if parallelize_prior is True."
)
self._aspire_instance = value
def __enter__(self):
self.original_log_likelihood = self.aspire_instance.log_likelihood
self.original_log_prior = self.aspire_instance.log_prior
if self.pool is not None:
logger.debug("Updating map function in log-likelihood method")
self.aspire_instance.log_likelihood = partial(
self.original_log_likelihood, map_fn=self.pool.map
)
if self.parallelize_prior:
logger.debug("Updating map function in log-prior method")
self.aspire_instance.log_prior = partial(
self.original_log_prior, map_fn=self.pool.map
)
return self.pool
def __exit__(self, exc_type, exc_value, traceback):
self.aspire_instance.log_likelihood = self.original_log_likelihood
self.aspire_instance.log_prior = self.original_log_prior
if self.close_pool:
logger.debug("Closing pool")
self.pool.close()
self.pool.join()
else:
logger.debug("Not closing pool")
[docs]
def logit(x: Array, eps: float | None = None) -> tuple[Array, Array]:
"""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.
"""
xp = array_namespace(x)
if eps:
x = xp.clip(x, eps, 1 - eps)
y = xp.log(x) - xp.log1p(-x)
log_j = (-xp.log(x) - xp.log1p(-x)).sum(-1)
return y, log_j
[docs]
def sigmoid(x: Array, eps: float | None = None) -> tuple[Array, Array]:
"""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.
"""
xp = array_namespace(x)
x = xp.divide(1, 1 + xp.exp(-x))
if eps:
x = xp.clip(x, eps, 1 - eps)
log_j = (xp.log(x) + xp.log1p(-x)).sum(-1)
return x, log_j
[docs]
def logsumexp(x: Array, axis: int | None = None) -> Array:
"""Implementation of logsumexp that works with array api.
This will be removed once the implementation in scipy is compatible.
"""
xp = array_namespace(x)
c = x.max()
return c + xp.log(xp.sum(xp.exp(x - c), axis=axis))
[docs]
def to_numpy(x: Array, **kwargs) -> np.ndarray:
"""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.
"""
if is_torch_array(x):
x = x.detach()
try:
return np.asarray(to_device(x, "cpu"), **kwargs)
except (ValueError, NotImplementedError, AttributeError):
return np.asarray(x, **kwargs)
[docs]
def asarray(x, xp: Any = None, dtype: Any | None = None, **kwargs) -> Array:
"""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.
"""
if dtype is not None:
kwargs["dtype"] = resolve_dtype(dtype, xp=xp)
# Handle DLPack conversion from JAX to PyTorch to avoid shape issues when
# passing JAX arrays directly to torch.asarray.
if is_jax_array(x) and is_torch_namespace(xp):
tensor = xp.utils.dlpack.from_dlpack(x)
if dtype is not None:
tensor = tensor.to(kwargs["dtype"])
return tensor
# Handle DLPack conversion from PyTorch to JAX to avoid issues with
# detaching tensors
if is_torch_array(x) and is_jax_namespace(xp):
import jax.dlpack
arr = jax.dlpack.from_dlpack(x.detach())
if dtype is not None:
arr = arr.astype(kwargs["dtype"])
return arr
if is_numpy_namespace(xp):
return to_numpy(x, **kwargs)
return xp.asarray(x, **kwargs)
[docs]
def determine_backend_name(
x: Array | None = None, xp: Any | None = None
) -> str:
"""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
-------
str
The name of the backend. If the backend cannot be determined, returns "unknown".
"""
if x is not None:
xp = array_namespace(x)
if xp is None:
raise ValueError(
"Either x or xp must be provided to determine backend."
)
for name, is_namespace_fn in IS_NAMESPACE_FUNCTIONS.items():
if is_namespace_fn(xp):
return name
return "unknown"
[docs]
def resolve_dtype(dtype: Any | str | None, xp: Any) -> Any | None:
"""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
-------
Any | None
The resolved dtype object compatible with ``xp`` (or None if unspecified).
"""
if dtype is None or xp is None:
return dtype
if isinstance(dtype, str):
dtype_name = _dtype_to_name(dtype)
if is_torch_namespace(xp):
resolved = getattr(xp, dtype_name, None)
if resolved is None:
raise ValueError(
f"Unknown dtype '{dtype}' for namespace {xp.__name__}"
)
return resolved
try:
return xp.dtype(dtype_name)
except (AttributeError, TypeError, ValueError):
resolved = getattr(xp, dtype_name, None)
if resolved is not None:
return resolved
raise ValueError(
f"Unknown dtype '{dtype}' for namespace {getattr(xp, '__name__', xp)}"
)
if is_torch_namespace(xp):
return dtype
try:
return xp.dtype(dtype)
except (AttributeError, TypeError, ValueError):
return dtype
def _dtype_to_name(dtype: Any | str | None) -> str | None:
"""Extract a canonical (lowercase) name for a dtype-like object."""
if dtype is None:
return None
if isinstance(dtype, str):
name = dtype
elif hasattr(dtype, "name") and getattr(dtype, "name"):
name = dtype.name
elif hasattr(dtype, "__name__"):
name = dtype.__name__
else:
text = str(dtype)
if text.startswith("<class '") and text.endswith("'>"):
text = text.split("'")[1]
if text.startswith("dtype(") and text.endswith(")"):
inner = text[6:-1].strip("'\" ")
text = inner or text
name = text
name = name.split(".")[-1]
return name.strip(" '\"<>").lower()
[docs]
def convert_dtype(
dtype: Any | str | None,
target_xp: Any,
*,
source_xp: Any | None = None,
) -> Any | None:
"""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
-------
Any | None
The dtype object compatible with ``target_xp`` (or None if ``dtype`` is None).
"""
if dtype is None:
return None
if target_xp is None:
raise ValueError("target_xp must be provided to convert dtype.")
target_name = getattr(target_xp, "__name__", "")
dtype_module = getattr(dtype, "__module__", "")
if dtype_module.startswith(target_name):
return dtype
if is_torch_namespace(target_xp) and str(dtype).startswith("torch."):
return dtype
name = _dtype_to_name(dtype)
if not name:
raise ValueError(f"Could not infer dtype name from {dtype!r}")
candidates = dict.fromkeys(
[name, name.lower(), name.upper(), name.capitalize()]
)
last_error: Exception | None = None
for candidate in candidates:
try:
return resolve_dtype(candidate, target_xp)
except ValueError as exc:
last_error = exc
# Fallback to direct attribute lookup
attr = getattr(target_xp, name, None) or getattr(
target_xp, name.lower(), None
)
if attr is not None:
return attr
raise ValueError(
f"Unable to convert dtype {dtype!r} to namespace {target_name}"
) from last_error
[docs]
def copy_array(x, xp: Any = None) -> Array:
"""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
-------
Array
The copied array.
"""
if xp is None:
xp = array_namespace(x)
# torch does not play nicely since it complains about copying tensors
if is_torch_namespace(xp):
if is_torch_array(x):
return xp.clone(x)
else:
return xp.as_tensor(x)
else:
try:
return xp.copy(x)
except (AttributeError, TypeError):
# Fallback for array APIs that do not have a copy method
return xp.array(x, copy=True)
[docs]
def effective_sample_size(log_w: Array) -> float:
xp = array_namespace(log_w)
return xp.exp(xp.asarray(logsumexp(log_w) * 2 - logsumexp(log_w * 2)))
@contextmanager
[docs]
def disable_gradients(xp, inference: bool = True):
"""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.
"""
if is_torch_namespace(xp):
if inference:
with xp.inference_mode():
yield
else:
with xp.no_grad():
yield
else:
yield
[docs]
def encode_dtype(xp, dtype):
"""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
-------
str
The encoded dtype.
"""
if dtype is None:
return None
return {
"__dtype__": True,
"xp": xp.__name__,
"dtype": _dtype_to_name(dtype),
}
[docs]
def decode_dtype(xp, encoded_dtype):
"""Decode a dtype from an HDF5 file.
Parameters
----------
xp : module
The array API module to use.
encoded_dtype : dict
The encoded dtype.
Returns
-------
dtype
The decoded dtype.
"""
if isinstance(encoded_dtype, dict) and encoded_dtype.get("__dtype__"):
if encoded_dtype["xp"] != xp.__name__:
raise ValueError(
f"Encoded dtype xp {encoded_dtype['xp']} does not match "
f"current xp {xp.__name__}"
)
if is_torch_namespace(xp):
return getattr(xp, encoded_dtype["dtype"].split(".")[-1])
else:
return xp.dtype(encoded_dtype["dtype"].split(".")[-1])
else:
return encoded_dtype
[docs]
def encode_samples(samples: BaseSamples) -> dict:
"""Encode a BaseSamples object for storage in an HDF5 file.
Parameters
----------
samples : BaseSamples
The samples to encode.
Returns
-------
dict
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.
"""
dictionary = samples._encode_for_hdf5()
return {
"__samples__": True,
"samples_type": type(samples).__name__,
"data": dictionary,
}
[docs]
def decode_samples(encoded_samples: dict) -> BaseSamples:
"""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 :code:`encode_samples` function.
Raises
------
ValueError
If the encoded_samples dictionary does not have the expected format or
if the samples type is unknown.
"""
if isinstance(encoded_samples, dict) and encoded_samples.get(
"__samples__"
):
samples_type = encoded_samples["samples_type"]
data = encoded_samples["data"]
# TODO: this could support user-defined samples
samples_module = importlib.import_module("aspire.samples")
SamplesClass = getattr(samples_module, samples_type, None)
if SamplesClass is None:
raise ValueError(
f"Unknown samples type '{samples_type}' in encoded data"
)
return SamplesClass._decode_from_dictionary(data)
else:
raise ValueError("Invalid encoded samples format")
[docs]
def encode_for_hdf5(value: Any) -> Any:
"""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`
"""
from .samples import BaseSamples
if is_jax_array(value) or is_torch_array(value):
return to_numpy(value)
if isinstance(value, BaseSamples):
value = encode_samples(value)
if isinstance(value, CallHistory):
return value.to_dict(list_to_dict=True)
if isinstance(value, np.ndarray):
return value
if isinstance(value, Path):
value = str(value)
if isinstance(value, (int, float, str)):
return value
if isinstance(value, (list, tuple)):
if all(isinstance(v, str) for v in value):
dt = h5py.string_dtype(encoding="utf-8")
return np.array(value, dtype=dt)
return [encode_for_hdf5(v) for v in value]
if isinstance(value, set):
return {encode_for_hdf5(v) for v in value}
if isinstance(value, dict):
if not value:
return "__empty_dict__"
else:
return {k: encode_for_hdf5(v) for k, v in value.items()}
if value is None:
return "__none__"
return value
[docs]
def decode_from_hdf5(value: Any) -> Any:
"""Decode a value loaded from an HDF5 file, reversing encode_for_hdf5."""
if isinstance(value, bytes): # HDF5 may store strings as bytes
value = value.decode("utf-8")
if isinstance(value, str):
if value == "__none__":
return None
if value == "__empty_dict__":
return {}
return value
if isinstance(value, np.ndarray):
# Try to collapse 0-D arrays into scalars
if value.shape == ():
return value.item()
if value.dtype.kind in {"S", "O", "U"}:
try:
return value.astype(str).tolist()
except Exception:
# fallback: leave as ndarray
return value
return value
if isinstance(value, list):
return [decode_from_hdf5(v) for v in value]
if isinstance(value, tuple):
return tuple(decode_from_hdf5(v) for v in value)
if isinstance(value, set):
return {decode_from_hdf5(v) for v in value}
if isinstance(value, dict):
if "__samples__" in value:
return decode_samples(value)
else:
return {
k.decode("utf-8"): decode_from_hdf5(v)
for k, v in value.items()
}
# Fallback for ints, floats, strs, etc.
return value
[docs]
def dump_pickle_to_hdf(memfp, fp, path=None, dsetname="state"):
"""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".
"""
memfp.seek(0)
bdata = np.frombuffer(memfp.read(), dtype="S1")
target = fp.require_group(path) if path is not None else fp
if dsetname not in target:
target.create_dataset(
dsetname, shape=bdata.shape, maxshape=(None,), dtype=bdata.dtype
)
elif bdata.size != target[dsetname].shape[0]:
target[dsetname].resize((bdata.size,))
target[dsetname][:] = bdata
[docs]
def dump_state(
state,
fp,
path=None,
dsetname="state",
protocol=pickle.HIGHEST_PROTOCOL,
):
"""Pickle a state object and store it in an HDF5 dataset."""
memfp = BytesIO()
pickle.dump(state, memfp, protocol=protocol)
dump_pickle_to_hdf(memfp, fp, path=path, dsetname=dsetname)
[docs]
def resolve_xp(xp_name: str | None):
"""
Resolve a backend name to the corresponding array_api_compat module.
Returns None if the name is None or cannot be resolved.
"""
if xp_name is None:
return None
name = xp_name.lower()
if name.startswith("array_api_compat."):
name = name.removeprefix("array_api_compat.")
try:
if name in {"numpy", "numpy.ndarray"}:
import array_api_compat.numpy as np_xp
return np_xp
if name in {"jax", "jax.numpy"}:
import jax.numpy as jnp
return jnp
if name in {"torch"}:
import array_api_compat.torch as torch_xp
return torch_xp
except Exception:
logger.warning(
"Failed to resolve xp '%s', defaulting to None", xp_name
)
return None
[docs]
def 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.
"""
if xp is None or is_numpy_namespace(xp) or is_jax_namespace(xp):
return None
try:
from array_api_compat import device
return device(x)
except Exception:
return None
[docs]
def 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.
"""
if device is None:
return x
if xp is None or is_numpy_namespace(xp) or is_jax_namespace(xp):
return x
try:
return to_device(x, device)
except Exception:
logger.warning(
"Failed to move array to device %s; leaving on current device",
device,
)
return x
[docs]
def recursively_save_to_h5_file(h5_file, path, dictionary):
"""Save a dictionary to an HDF5 file with flattened keys under a given group path."""
# Ensure the group exists (or open it if already present)
group = h5_file.require_group(path)
def _save_flattened(g, prefix, d):
for key, value in d.items():
full_key = f"{prefix}.{key}" if prefix else key
if isinstance(value, dict):
_save_flattened(g, full_key, value)
else:
try:
g.create_dataset(full_key, data=encode_for_hdf5(value))
except (TypeError, ValueError) as error:
logger.warning(
f"Failed to save key {full_key} with value {value} to HDF5 file: {error}"
)
logger.warning(
"Attempting to save as a string representation instead."
)
try:
# Try saving as a string
dt = h5py.string_dtype(encoding="utf-8")
g.create_dataset(
full_key, data=np.array(str(value), dtype=dt)
)
except Exception:
raise RuntimeError(
f"Cannot save key {full_key} with value {value} to HDF5 file."
) from error
_save_flattened(group, "", dictionary)
[docs]
def load_from_h5_file(h5_file, path):
"""Load a flattened dictionary from an HDF5 group and rebuild nesting."""
group = h5_file[path]
result = {}
for key, dataset in group.items():
parts = key.split(".")
d = result
for part in parts[:-1]:
d = d.setdefault(part, {})
d[parts[-1]] = decode_from_hdf5(dataset[()])
return result
[docs]
def get_package_version(package_name: str) -> str:
"""Get the version of a package.
Parameters
----------
package_name : str
The name of the package.
Returns
-------
str
The version of the package.
"""
try:
module = __import__(package_name)
return module.__version__
except ImportError:
return "not installed"
[docs]
class AspireFile(h5py.File):
"""A subclass of h5py.File that adds metadata to the file."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._set_aspire_metadata()
def _set_aspire_metadata(self):
from . import __version__ as aspire_version
if self.mode in {"w", "w-", "a", "r+"}:
self.attrs["aspire_version"] = aspire_version
else:
aspire_version = self.attrs.get("aspire_version", "unknown")
if aspire_version != "unknown":
logger.warning(
f"Opened Aspire file created with version {aspire_version}. "
f"Current version is {aspire_version}."
)
[docs]
def update_at_indices(x: Array, slc: Array, y: Array) -> Array:
"""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
-------
Array
The updated array.
"""
warnings.warn(
"update_at_indices is deprecated and will be removed in a future version. Please use array-api-extra.at instead",
UserWarning,
)
return xpx.at(x, slc).set(y)
@dataclass
[docs]
class CallHistory:
"""Class to store the history of calls to a function.
Attributes
----------
args : list[tuple]
The positional arguments of each call.
kwargs : list[dict]
The keyword arguments of each call.
"""
[docs]
kwargs: list[dict]
[docs]
def to_dict(self, list_to_dict: bool = False) -> dict[str, Any]:
"""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.
"""
if list_to_dict:
return {
"args": {str(i): v for i, v in enumerate(self.args)},
"kwargs": {str(i): v for i, v in enumerate(self.kwargs)},
}
else:
return {
"args": [list(arg) for arg in self.args],
"kwargs": [dict(kwarg) for kwarg in self.kwargs],
}
[docs]
def 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.
"""
@wrapt.decorator
def wrapper(wrapped_func, instance, args, kwargs):
# If instance is provided, we're dealing with a method.
if instance:
# Attach `calls` attribute to the method's `__func__`, which is the original function
if not hasattr(wrapped_func.__func__, "calls"):
wrapped_func.__func__.calls = CallHistory([], [])
wrapped_func.__func__.calls.args.append(args)
wrapped_func.__func__.calls.kwargs.append(kwargs)
else:
# For standalone functions, attach `calls` directly to the function
if not hasattr(wrapped_func, "calls"):
wrapped_func.calls = CallHistory([], [])
wrapped_func.calls.args.append(args)
wrapped_func.calls.kwargs.append(kwargs)
# Call the original wrapped function
return wrapped_func(*args, **kwargs)
return wrapper(wrapped) if wrapped else wrapper
[docs]
def function_id(fn: Any) -> str:
"""Get a unique identifier for a function.
Parameters
----------
fn : Any
The function to get the identifier for.
Returns
-------
str
The unique identifier for the function.
"""
if isinstance(fn, functools.partial):
base = fn.func
else:
base = fn
return f"{base.__module__}:{getattr(base, '__qualname__', type(base).__name__)}"
[docs]
def enable_scipy_array_api() -> None:
"""Set the environment variable to enable array API compatibility in SciPy.
Set :code:`SCIPY_ARRAY_API=1` if it is not already set.
"""
import os
val = os.environ.get("SCIPY_ARRAY_API")
if val is None:
warnings.warn(
"SCIPY_ARRAY_API environment variable is not set. Setting it to '1'",
ImportWarning,
)
os.environ["SCIPY_ARRAY_API"] = "1"
elif val != "1":
warnings.warn(
f"SCIPY_ARRAY_API={val}, which may cause issues with "
"array API compatibility. Consider setting it to '1' to enable "
"array API support in SciPy.",
RuntimeWarning,
)