bilby.compat.utils.array_module

bilby.compat.utils.array_module(arr)[source]

Infer the array module (namespace) from the input argument. This is a generalization of the array_api_compat.array_namespace function that can handle a wider variety of inputs, including some nested structures.

This function determines which array library backend is being used by inspecting the input argument. It handles various input types and fallback mechanisms to ensure a valid array module is always returned.

The inference logic proceeds as follows: 1. If a single-element tuple is provided, extract the element first. 2. Attempt to use the array_api_compat.array_namespace() function

directly, which works for most array-like objects.

  1. If that fails, handle special cases: - Dictionaries: extract values and infer from non-string values - Builtin iterables (list, tuple, etc.): infer from elements - Builtin scalars: default to numpy - Pandas objects: default to numpy (treated as numpy-compatible) - Unknown types: log a warning and default to numpy

This is a best-effort function, but will not cover all possible edge cases.

If support for the general array API is not activated via BILBY_ARRAY_API=1, this always returns numpy.

Parameters:
arr: array-like, tuple, dict, or other type

The input argument to infer the array module from. Can be: - An array object (numpy, cupy, jax.numpy, etc.) - A tuple of arrays (single-element unwrapped) - A dictionary with array values - An iterable containing arrays - A builtin scalar or type

Returns:
module

The array namespace module (e.g., numpy, cupy, jax.numpy, etc.). Defaults to numpy if the module cannot be determined.

Examples

>>> import numpy as np
>>> import jax.numpy as jnp
>>> array_module(np.array([1, 2, 3]))
<module 'numpy' ...>

>>> array_module(jnp.array([1, 2, 3]))
<module 'jax.numpy' ...>

>>> array_module({'data': np.array([1, 2, 3])})
<module 'numpy' ...>

>>> array_module([np.array([1]), np.array([2])])
<module 'numpy' ...>

>>> array_module([1, jnp.array([2])])
<module 'jax.numpy' ...>

>>> array_module(5)
<module 'numpy' ...>