PR: Implement support for *Python Array API Standard*.

[thomasmansencal]

Summary

This PR implements support for the Python Array API Standard, enabling computations to dispatch onto alternative array backends:

• NumPy (default)
• JAX
• PyTorch (including Apple MPS)

Dispatch is currently opt-in and NumPy-only behaviour is unchanged by default. Once enabled, the backend is selected from the type of the input array. It can be enabled three ways:

1. Environment variable, set before importing Colour:

import os

os.environ["COLOUR_SCIENCE__ARRAY_API"] = "1"

import colour
import jax.numpy as jnp

colour.XYZ_to_sRGB(jnp.array([0.20654008, 0.12197225, 0.05136952]))
# Array([0.7057394 , 0.19248262, 0.2235417 ], dtype=float32)

2. Programmatically, toggle the global state at runtime:

import colour
import torch
from colour.utilities import set_array_api_enabled

set_array_api_enabled(True)

colour.XYZ_to_sRGB(torch.tensor([0.20654008, 0.12197225, 0.05136952]))
# tensor([0.7057, 0.1925, 0.2235], dtype=torch.float64)

3. Scoped context manager (also usable as a decorator), enable for a block only:

import colour
import torch
from colour.utilities import array_api_enable

with array_api_enable(True):
    colour.XYZ_to_sRGB(
        torch.tensor([0.20654008, 0.12197225, 0.05136952], device="mps")
    )
# tensor([0.7057, 0.1925, 0.2235], device='mps:0')

What's added

• Namespace-aware boundary helpers + a full xp_* operation surface in colour.utilities:
  • Namespace resolution: array_namespace, is_numpy_namespace, is_non_ndarray, trace_array_namespace
  • Boundary conversion: as_ndarray, cast_non_ndarray, xp_as_array / xp_as_float_array / xp_as_int_array, xp_astype, xp_ascontiguousarray
  • Shape & manipulation: xp_reshape, xp_squeeze, xp_atleast_1d / xp_atleast_2d, xp_broadcast_to, xp_matrix_transpose, xp_resize, xp_pad, xp_insert
  • Reductions & statistics: xp_average, xp_median, xp_nanmean, xp_trapezoid, xp_gradient
  • Element-wise math: xp_degrees / xp_radians, xp_sinc, xp_round, xp_nan_to_num
  • Linear algebra: xp_lstsq, xp_eig / xp_eigh, xp_create_diagonal
  • Sampling, interpolation & set operations: xp_linspace, xp_interp, xp_select, xp_isin, xp_setxor1d, xp_unique
  • Comparison & testing: xp_isclose, xp_assert_close, xp_assert_equal
• contextvars-backed global state (Array API enablement, domain-range scale, ndarray copy, caching) for thread/async safety.
• SciPy-free, dispatchable kernels replacing solver/interpolator hotspots: correlated colour temperature Gauss-Newton (colour.temperature.common), Jakob and Hanika (2019) trilinear interpolation, etc.
• Default complex precision: COLOUR_SCIENCE__DEFAULT_COMPLEX_DTYPE / set_default_complex_dtype.
• New public API: CIE_illuminant_D_series, msds_CIE_illuminant_D_series, msds_blackbody, msds_rayleigh_jeans.
• Cross-backend testing: an xp pytest fixture parametrising numpy/jax/torch/torch-mps, with mps_tolerance_absolute and mps_xfail markers for float32 precision, plus a cross-backend benchmark suite (utilities/benchmark.py).
• Documentation: a dedicated Array API Support section in advanced.rst.

Performance

Per-suite speed-up vs NumPy (best-of-3, HD inputs): speed-up = NumPy ÷ backend over cases succeeding on both, so higher = faster (e.g. 3.0× = 3× faster than NumPy; < 1.0× = slower). numpy (ms) is the summed best-of-3 over the suite's cases.

Suite	cases	numpy (ms)	jax	torch-cpu	torch-mps
conversion_graph	207	32973.5	3.1×	3.4×	18×
conversion_graph_iterative	4	38181.9	3.3×	1.8×	12×
difference	17	2230.7	2.4×	2.8×	30×
integration_array	2	211.3	2.6×	1.02×	15×
integration_object	6	3.8	1.3×	0.91×	0.90×
transfer_function	114	3892.1	3.0×	2.2×	18×
adaptation	7	1095.8	3.6×	4.3×	6.0×
characterisation	3	262.2	2.7×	3.7×	9.5×
recovery_array	4	1938.3	2.3×	2.7×	13×
recovery_object	3	508.8	0.60×	0.59×	0.50×
quality_array	4	538.3	2.3×	2.8×	3.3×
quality_object	5	8.1	0.13×	0.56×	0.07×
volume	2	37.5	0.92×	0.93×	1.1×
volume_iterative	2	1957.4	1.03×	1.03×	16×
phenomena	5	176.6	2.2×	3.7×	25×
temperature_array	4	339.3	3.0×	2.5×	23×
temperature_iterative	4	1177.8	0.84×	1.6×	0.74×
blindness	3	166.0	5.4×	16×	7.7×
contrast	1	82.2	2.8×	3.4×	24×
generators_array	3	61.0	2.1×	1.9×	3.3×
generators_object	6	11.5	0.75×	0.67×	0.56×
photometry	3	0.1	0.11×	0.22×	0.03×
overall	409	85854.3	2.8×	2.2×	11×

NumPy is the baseline (1.00×). JAX dispatches asynchronously, so every timed operation is synchronised with jax.block_until_ready before the clock stops: the jax column measures completed computation, not enqueue latency.

Measured on an Apple M1 Max (10-core, 32 GB), macOS 15.7, Python 3.13, NumPy 2.3, PyTorch 2.9, JAX 0.8; 409 cases across 22 suites.

Preflight

Code Style and Quality

• Unit tests have been implemented and passed.
• Pyright static checking has been run and passed.
• Pre-commit hooks have been run and passed.
• [N/A] New transformations have been added to the Automatic Colour Conversion Graph.
• New transformations have been exported to the relevant namespaces, e.g. colour, colour.models.

Documentation

• New features are documented along with examples if relevant.
• The documentation is Sphinx and numpydoc compliant.

[MichaelMauderer]

Benchmarks on my system with AMD Ryzen Threadripper 2950X / 64 GiB RAM / Radeon RX 7900 XT.

Suite	cases	numpy (ms)	jax	torch-cpu	torch-cuda
conversion_graph	207	61714.0	2.58×	4.07×	22.10×
conversion_graph_iterative	4	75222.7	2.72×	2.62×	51.60×
difference	17	3882.5	2.10×	3.19×	87.89×
integration_array	2	528.5	3.61×	3.72×	11.12×
integration_object	6	10.7	1.47×	0.44×	3.69×
transfer_function	114	6447.1	2.29×	3.56×	19.23×
adaptation	7	1987.7	2.54×	4.43×	6.27×
characterisation	3	468.5	2.11×	4.25×	8.20×
recovery_array	4	2778.4	1.83×	2.56×	103.89×
recovery_object	3	1678.7	0.61×	0.60×	0.57×
quality_array	4	899.8	1.88×	4.27×	5.94×
quality_object	5	13.5	0.10×	0.39×	0.27×
volume	2	71.3	0.96×	0.98×	0.97×
volume_iterative	2	5116.0	0.99×	0.99×	1.01×
phenomena	5	273.9	2.64×	5.33×	52.90×
temperature_array	4	481.7	2.64×	2.77×	87.17×
temperature_iterative	4	1552.6	0.36×	0.86×	1.12×
blindness	3	350.4	2.64×	11.91×	7.49×
contrast	1	136.7	2.72×	4.64×	75.92×
generators_array	3	85.2	2.03×	2.25×	4.83×
generators_object	6	24.2	0.65×	1.35×	1.60×
photometry	3	0.2	0.12×	0.22×	0.11×
all	409	163724.2	2.26×	2.77×	11.04×

Looks very much in line with your results @thomasmansencal .

Do we need dynamic opt-in for the API? I wonder if it is enough to have the runtime check for the available backends. The opt-in seems mostly useful for backwards compatibility, but I think we need to raise our version number for these changes anyway.

I'll do a more thorough review tomorrow.

[lassefschmidt]

First of all : great job !

I had a look at this from an autodiff angle (trying to see if computational graph breaks even though it shouldn't). As in trying to see where we can compute gradients based on function inputs and where not -- this would be a huge unlock for everyone working on inverse models.

More than happy to propose fixes for the below mentioned breaks in the computational graph, if you are willing to include this in the scope for our upcoming release.

• PR head: b6e63a904d75ab46cb0c64fe3dbaee4ee7cfb4e2
• Comparison develop: 907242acd5e514a94b626a9dcf4bfe953aa0b8dc
• PyTorch 2.13.0, CPU, float64, Array API dispatch enabled

Full evaluation details available here : drilldown.md

Accidental graph breaks to fix, in priority order

Priority	Location	Current break	Recommended direction
P0	colour.difference.delta_e: CIE 2000 and HyCH	dataclasses.astuple(...) deep-copies every tensor field and fails for non-leaf Torch tensors.	Shallow-unpack attributes directly, or use tuple(getattr(spec, field.name) for field in dataclasses.fields(spec)). Add direct and spectral-chain gradient tests.
P0	colour.continuous.AbstractContinuousFunction.copy / colour.colorimetry.reshape_sd and reshape_msds	deepcopy(self) fails on non-leaf Torch tensors. Re-promotion after SciPy would also restore values only after losing history.	Implement a backend-aware structural copy and backend-native fixed-grid interpolation/resampling. Do not convert a grad-tracking tensor to NumPy and back.
P0	colour.recovery.XYZ_to_sd_Jakob2019	Explicit as_ndarray(XYZ) at the public boundary.	Keep XYZ and iterative coefficients in the selected namespace; avoid Python-float convergence outputs in the differentiable value path.
P0/P1	colour.recovery.XYZ_to_sd_Meng2015	Explicit NumPy conversion plus scipy.optimize.minimize.	Either document this method as non-autodiff, provide an Array API solver, or implement a custom/implicit backward. Do not silently return a detached spectrum for Torch input.
P1	colour.quality.spectral_similarity_index	SciPy convolve1d(as_ndarray(...)) detaches even if reshape is fixed.	Express the fixed three-tap convolution as backend-native padding/convolution or a fixed matrix multiplication. Preserve round_result=False as the differentiable mode.
P1	colour.difference.sd_to_metamerism_index	Namespace is inferred from NumPy A_r, then NumPy consumes Torch spectral values.	Resolve namespace from spectral values and weights together, then promote A_r/A_t to the spectrum backend/device.
P1	colour.colorimetry.photometry / xp_trapezoid	A NumPy wavelength axis makes Torch's trapezoid call fail and triggers a NumPy fallback that detaches spectral values.	Promote x to the selected namespace/device before the native call; use dx for fixed uniform shapes where appropriate. Add ordinary Colour-SD gradient tests.
P1	colour.colorimetry.dominant	Dominant/complementary wavelength is discrete as expected, but selected intersection coordinates and purity propagate NaN gradients.	Keep wavelength explicitly non-differentiable; compute valid candidate intersections without NaN-producing inactive branches and test finite local purity gradients inside a stable locus segment.
P1	colour.temperature.robertson1968	In-place /= and += operations modify tensors needed for backward.	Replace in-place normalisation/offsets with out-of-place expressions and test both uv→CCT and CCT/Duv→uv. Document interval selection as piecewise.
P1	Generic CCT_to_uv, Planck 1900 entry	Dispatcher pair contract does not match the scalar-CCT Planck implementation.	Adapt [CCT, Duv] explicitly or remove Planck from the generic pair mapping; keep the direct Planckian-locus API.
P1	sd_to_XYZ result cache	Cache retains an old autograd graph; a second backward through a cache hit fails after the first graph is freed.	Bypass output caching if any input/value has requires_grad=True, or scope caches to no-grad computations.
P1	CRI/CQS/TLCI/TLMF	High-level object reshaping blocks the whole otherwise mostly array-dispatched pipeline.	Fix the shared spectral copy/reshape boundary first, then retest. CIEDE2000 and discrete mask qualifications remain.
P2	Otsu/Smits recovery	Useful local gradients can be mistaken for global differentiability.	Document piecewise behaviour and add tests on stable branches plus explicit boundary tests. Consider smooth-selection alternatives only as optional optimisation methods.
P2	contrast.sigma_Barten1999	Torch d with default scalar constants creates mixed NumPy/Torch arguments to hypot.	Promote every scalar/default operand to the resolved namespace, not only a subset.

Suggested regression-test structure

The PR currently demonstrates cross-backend numerical execution. Autodiff needs separate assertions. A compact test helper should:

1. Create torch.float64 leaf tensors with requires_grad=True.
2. Run the Colour operation with Array API enabled and result caching disabled.
3. Assert the result is a Torch tensor, requires_grad is true, and grad_fn is present.
4. Call torch.autograd.grad(output.sum(), every_input) and assert every intended gradient is non-None and finite.
5. Run torch.autograd.gradcheck at regular points for smooth/local-piecewise kernels.
6. Keep separate tests for documented non-smooth boundaries; do not require gradcheck exactly at hue wraps, equal colours, clip thresholds, transfer-function cutoffs, cluster boundaries, or channel ties.
7. Add end-to-end spectral tests, not only isolated function tests:
   • raw spectra → XYZ → Lab → each Lab-based ΔE;
   • raw spectra → XYZ → ICtCp/CAM UCS → matching ΔE;
   • raw spectra → XYZ → every whiteness and yellowness method;
   • raw spectra → XYZ → uv → every CCT method, checking finite gradients rather than only grad_fn;
   • ordinary SpectralDistribution objects → luminous flux/efficiency/efficacy with NumPy wavelength axes;
   • raw spectra → XYZ → xy → purity, explicitly rejecting NaN gradients;
   • XYZ → each recovery method → raw spectrum integration → Lab loss;
   • spectral curve → CIE 2017 and TM-30 fidelity scores;
   • repeat a cached spectral computation across separate backward passes to prevent graph-reuse regressions.