RFC: wrap any array API standard compatible array

[lucascolley]

Preface: I'm not familiar with the history behind https://pint.readthedocs.io/en/stable/user/numpy.html#Array-Type-Support, but I have experience adding support for alternative array types to SciPy. For context, see https://numpy.org/neps/nep-0056-array-api-main-namespace.html#abstract, which superseded NEP 30.

---

The array API standard is being utilised in libraries like SciPy and scikit-learn to add support for array libraries like CuPy, PyTorch, and JAX. It seems feasible that Pint could wrap any library which is compatible with the standard.

The model for how this can be done is demonstrated in https://github.com/mdhaber/marray. Marray adds masked array support to standard-compatible arrays. Similar to Pint's Quantity, an MArray is some data together with a mask:

MArray(
    Array([1, 2], dtype=array_api_strict.int64),
    Array([False,  True], dtype=array_api_strict.bool)
)

Marray exposes one function, get_namespace, which takes as input a standard-compatible namespace, and returns a new standard-compatible namespace with the additional capabilities of creating and correctly processing masked arrays. Under the hood, this namespace wraps each function to handle anything specific to masks, and otherwise dispatches to the underlying array namespace. So for pint, pint.get_namespace(torch) would return a namespace for pint-wrapped PyTorch tensors.

MArrays also expose this namespace via the __array_namespace__ method, which allows libraries to operate on them like any other standard array:

# get the namespace for input pint(torch) array x
xp = x.__array_namespace__()
y = xp.exp(x) # returns another pint(torch) array
z = xp.asarray([1, 2, 3]) # also a pint(torch) array
...

Marray currently passes almost the entirety of https://github.com/data-apis/array-api-tests. It would be cool to see if Pint could do the same!

This would bring a few benefits:

• for users, they can use pint with whatever standard-compatible library they like
• for libraries which depend on pint, they can also then work on supporting any standard-compatible arrays
• users can pass Pint Quantitys to any library which supports standard-compatible arrays, like SciPy (eventually!)

---

@nabobalis I know you had expressed interest in this.

[andrewgsavage]

yes, there is interest. there have been a few issues relating to this in the past but it requires someone to start it! I suggest we create a seperate repo for this as it allows the array code to be iterated and released faster and independently to pint. maybe pint-array or pint-array-api?

[lucascolley]

alright! Extremely rough POC at https://github.com/lucascolley/pint-array:

In [1]: from pint import UnitRegistry

In [2]: from pint_array import get_namespace

In [3]: import array_api_strict as xp

In [4]: pxp = get_namespace(xp)

In [5]: ureg = UnitRegistry()

In [6]: x = pxp.asarray(xp.arange(6), units=ureg.meter)

In [7]: pxp.sum(x)
Out[7]: <Quantity(15, 'meter')>

In [9]: x = pxp.asarray(x, dtype=xp.float64)

In [10]: pxp.var(x)
Out[10]: <Quantity(2.9166666666666665, 'meter ** 2')>

In [11]: x = pxp.asarray(x, units=ureg.degC)

In [12]: pxp.sum(x)
---------------------------------------------------------------------------
OffsetUnitCalculusError                   Traceback (most recent call last)
Cell In[12], line 1
----> 1 pxp.sum(x)

File ~/programming/pint-array/pint_array/__init__.py:42, in get_namespace.<locals>.sum(x, axis, dtype, keepdims)
     40 units = x.units
     41 magnitude = xp.sum(x, axis=axis, dtype=dtype, keepdims=keepdims)
---> 42 units = (1 * units + 1 * units).units
     43 return UnitRegistry.Quantity(magnitude, units)

File ~/programming/pint-array/.pixi/envs/default/lib/python3.13/site-packages/pint/facets/plain/quantity.py:849, in PlainQuantity.__add__(self, other)
    846 if isinstance(other, datetime.datetime):
    847     return self.to_timedelta() + other
--> 849 return self._add_sub(other, operator.add)

File ~/programming/pint-array/.pixi/envs/default/lib/python3.13/site-packages/pint/facets/plain/quantity.py:101, in check_implemented.<locals>.wrapped(self, *args, **kwargs)
     99 elif isinstance(other, list) and other and isinstance(other[0], type(self)):
    100     return NotImplemented
--> 101 return f(self, *args, **kwargs)

File ~/programming/pint-array/.pixi/envs/default/lib/python3.13/site-packages/pint/facets/plain/quantity.py:825, in PlainQuantity._add_sub(self, other, op)
    823     units = other._units
    824 else:
--> 825     raise OffsetUnitCalculusError(self._units, other._units)
    827 return self.__class__(magnitude, units)

OffsetUnitCalculusError: Ambiguous operation with offset unit (degree_Celsius, degree_Celsius). See https://pint.readthedocs.io/en/stable/user/nonmult.html for guidance.

With a bit of work it should be relatively straightforward based on the work already done to wrap NumPy (I hope).

[andrewgsavage]

Thanks, that's very helpful to understand how it works.

Assuming the function signatures are the same for numpy asn the array API, it may be possible to change the implements decorator to also create functions for pint_namespace
That would be fine to do in pint.

[lucascolley]

Assuming the function signatures are the same for numpy asn the array API, it may be possible to change the implements decorator to also create functions for pint_namespace
That would be fine to do in pint.

Yes, it seems like we should be able to reuse a lot of that code. Given that NumpyWrapper supports functions (and ufuncs) which aren't in the array API standard, I think it will be easier to keep them separate for now.

I just pushed a change to the repo such that we can define an ArrayUnitQuantity, which is a pint.Quantity with additional methods from ArrayQuantity. Since the return value of the __array_namespace__ method is the pint_namespace which is itself defined using ArrayUnitQuantitys, I'm not sure yet how ArrayQuantity could become a standalone facet to integrate into pint.UnitRegistry. But maybe that isn't needed!

[mhvk]

xref astropy/astropy#17588

This is probably as good a moment as any to mention that @nstarman and I have been working on a new Quantity implementation that aims to be array API compatible and agnostic of the unit system - see https://github.com/astropy/quantity-2.0 and the design considerations at astropy/astropy-APEs#91 (especially the PDF linked on top).

As our hope in part is to replace astropy's Quantity, our minimal goal is to provide full coverage of the array API; it would be the full numpy API for ndarray (i.e., including out arguments, specializations of ufuncs like np.add.{at, reduce, reduceat, accumulate}, and gufuncs.

Anyway, the part that is perhaps most relevant here is the unit-agnostic part - this would need some standardization of how units interact with arrays, perhaps new dunder methods of some agreement on operators (e.g., astropy can attach a unit via value << unit - maybe that could become a standard?)

[nstarman]

A shared units & Quantity API would be fantastic!
I would vote for defining a quantity-api libraries with some runtime-checkable Protocols.
I mocked up something along these lines in https://github.com/nstarman/units/tree/main/src/units/api

[andrewgsavage]

@mhvk , when you say unit system, do you mean the unit module, pint or astropy.units etc?

the linked pdf refers to a Quantity API. If I am understanding it correctly, if pint, astropy.units and any other unit library followed it, then there could be a quantity-array library that works for both pint and astropy.units.

could then make pint-pandas and pint-xarray work for astropy and others too

[lucascolley]

Yep, that sounds perfect. The namespaces I'm building in pint-array just need an array library, a quantity class from which to inherit, and some units compatible with that quantity. Currently, that's pint.Quantity and pint.UnitRegistry, but it should be relatively simple to swap those out for a new standard if you folks can figure one out!

That's probably the hard part; in the meantime, I'll work on getting these array-quantity namespaces up to standard-compliance, while attempting to do the right thing with units. I could probably use some help writing tests for correct unit-handling at some point!

[mhvk]

@lucascolley - spurred on by this conversation, I've pushed my implementation of ufucs/element-wise operations to astropy/quantity-2.0#22 - tests/test_element_wise_functions.py may be fairly directly useful to you.

[andrewgsavage]

Yes, there aren't all that many classes and functions needed for the array API, based on pint's numpy array_function implementation. I think a Quantity API defining the classes and functions needed to write the pint-array library would be a good first version. Is https://github.com/astropy/quantity-2.0 the best place for such a discussion on a Quantity API?

There are a load of numpy tests that have been added as people implement numpy functions.

[lucascolley]

What do we think about creation-like functions such as ones_like? Should the result be dimensionless or have the same units as the input array? Pint currently returns a dimensionless array for NumPy arrays.

[andrewgsavage]

that was being changed to have the same units as the input array, but has got stuck #882

[lucascolley]

Up until this point everything has been working with pint.UnitRegistry. I'm now trying to implement https://data-apis.org/array-api/latest/API_specification/generated/array_api.concat.html#array_api.concat, my implementation of which makes use of https://pint.readthedocs.io/en/stable/api/facets.html#pint.facets.plain.PlainQuantity.m_as. That takes us through to this generic _convert function:

pint/pint/facets/plain/registry.py

Lines 1087 to 1097 in 74b7086

	# factor is type float and if our magnitude is type Decimal then
	# must first convert to Decimal before we can '*' the values
	if isinstance(value, Decimal):
	factor = Decimal(str(factor))
	elif isinstance(value, Fraction):
	factor = Fraction(str(factor))
	if inplace:
	value *= factor
	else:
	value = value * factor

In the array API standard:

Behavior is not specified when mixing a Python float and an array with an integer data type; this may give float32, float64, or raise an exception.

and so the multiplication here can throw an exception for integer arrays.

Do you think we could try to address this problem in the plain registry @andrewgsavage ? It will be a bit hacky without adding proper array API awareness to Pint, but I think something like

if hasattr(value, __array_namespace__) and value.__array_namespace__().isdtype(value.dtype, "integral"):
    factor = int(factor)

might fix the integer array case.

I suppose the alternative is for me to subclass UnitRegistry and override the _convert method. But it might be nice to avoid that if possible?