Skip to content

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

Open
KelSolaar wants to merge 1 commit into
developfrom
feature/array-api-support
Open

PR: Implement support for *Python Array API Standard*.#1406
KelSolaar wants to merge 1 commit into
developfrom
feature/array-api-support

Conversation

@KelSolaar

@KelSolaar KelSolaar commented Jun 2, 2026

Copy link
Copy Markdown
Member

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 31455.1 4.4× 3.5× 11×
conversion_graph_iterative 4 36874.8 3.3× 1.8× 12×
difference 17 2139.4 3.1× 3.0× 18×
integration_array 2 201.5 644× 1.00× 15×
integration_object 6 3.2 1.3× 0.77× 0.76×
transfer_function 114 3721.2 15× 2.3× 10×
adaptation 7 1066.4 4.4× 4.8× 6.0×
characterisation 3 249.9 6.4× 4.0× 9.5×
recovery_array 4 1878.3 2.7× 3.0× 7.2×
recovery_object 3 452.6 0.65× 0.66× 0.48×
quality_array 4 488.9 2.4× 2.8× 1.5×
quality_object 5 5.1 0.10× 0.39× 0.05×
volume 2 39.5 1.1× 0.98× 1.2×
volume_iterative 2 1856.4 1.00× 0.97× 16×
phenomena 5 172.8 3.0× 4.5× 27×
temperature_array 4 294.7 4.7× 1.6× 18×
temperature_iterative 4 1015.5 0.75× 1.3× 0.53×
blindness 3 163.4 271× 15× 8.1×
contrast 1 78.6 5.1× 3.5× 26×
generators_array 3 51.8 3.3× 1.7× 1.9×
generators_object 6 12.7 1.3× 1.3× 0.84×
photometry 3 0.1 0.08× 0.18× 0.02×
overall 409 82221.8 3.3× 2.2× 8.5×

NumPy is the baseline (1.00×). 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.

@KelSolaar KelSolaar force-pushed the feature/array-api-support branch 4 times, most recently from bc223a6 to 3163508 Compare June 4, 2026 11:24
@KelSolaar KelSolaar changed the title Implement support for *Python Array API Standard*. PR: Implement support for *Python Array API Standard*. Jun 6, 2026
@KelSolaar KelSolaar force-pushed the feature/array-api-support branch 21 times, most recently from 98a7bdf to d825253 Compare June 13, 2026 21:05
*Colour* now dispatches array operations to the caller's backend (*NumPy*,
*JAX*, *PyTorch*) through the array-namespace machinery in
`colour.utilities.array`. Beyond the mechanical *NumPy* to namespace
conversion, this commit bundles the behaviour and public API changes
documented below so that they remain discoverable under `git blame` and
`git bisect`.

## Features

- Support for the *Python Array API Standard* was implemented: array
  operations dispatch to the input backend (*NumPy*, *JAX*, *PyTorch*)
  through the new `colour.utilities.array_namespace` and `xp_*` utilities,
  toggled with `colour.utilities.is_array_api_enabled` and
  `colour.utilities.set_array_api_enabled`.
- `colour.utilities.is_array_api_compat_installed` and
  `colour.utilities.is_array_api_extra_installed` were added.
- `colour.colorimetry.interpolate_signal`,
  `colour.colorimetry.extrapolate_signal` and
  `colour.colorimetry.trim_signal` were added, sharing the spectral
  distribution and multi-spectral distributions resampling implementation.
- `colour.colorimetry.msds_blackbody`,
  `colour.colorimetry.msds_rayleigh_jeans`,
  `colour.colorimetry.CIE_illuminant_D_series`,
  `colour.colorimetry.msds_CIE_illuminant_D_series` and
  `colour.colorimetry.msds_to_XYZ_tristimulus_weighting_factors_ASTME308`
  were added.
- `colour.appearance.eccentricity_factor_Hellwig2022` and
  `colour.appearance.hue_angle_dependency_Hellwig2022` were added.
- `colour.appearance.XYZ_to_Nayatani95` now computes the hue quadrature
  `H` correlate, previously left unset.

## Performance

- The multi-spectral distributions paths of `colour.colour_fidelity_index`,
  `colour.colour_quality_scale` and `colour.colour_rendering_index` were
  vectorised.
- The `colour.temperature` correlated colour temperature solvers were
  vectorised, replacing the *SciPy* `minimize` calls with closed-form
  Gauss-Newton iterations.

## Fixes

- `colour.colour_rendering_index`: the *"CIE 2024"* `Q_a` general index now
  averages test colour samples 1 to 8, it was averaging all 15.
- `colour.adaptation.chromatic_adaptation_Li2025` now applies domain and
  range scaling.
- The `COLOUR_SCIENCE__FILTER_COLOUR_WARNINGS` environment variable is now
  honoured correctly.

## Changes

- `colour.utilities.set_caching_enable`,
  `colour.utilities.set_ndarray_copy_enable` and
  `colour.algebra.set_spow_enable` were renamed to `set_caching_enabled`,
  `set_ndarray_copy_enabled` and `set_spow_enabled` respectively, without
  aliases.
- *Multiprocessing* support was removed: `disable_multiprocessing`,
  `multiprocessing_pool` and `ParallelForMultiprocess`.
- Around 80 internal appearance helpers were removed from the
  `colour.appearance` modules `__all__` (`ciecam02`, `ciecam16`,
  `hellwig2022`, `hunt`, `nayatani95`, `llab`, `atd95`).
- `colour.quality.cfi2017.sd_reference_illuminant` and
  `colour.quality.cfi2017.CCT_reference_illuminant` were removed, orphaned
  by the vectorised reference illuminant path.
- The appearance models `compute_H` argument now defaults to `False`.
- The `*_to_msds` definitions now return a `MultiSpectralDistributions`
  instance by default instead of a `numpy.ndarray`.
- `colour.algebra.least_square_mapping_MoorePenrose` now uses batched,
  greater than 2-D, matrix multiplication semantics.
- The *Jiang et al. (2013)* principal component analysis dropped its
  covariance-matrix path; its reference basis functions were regenerated.
- The `colour.temperature` solvers reference values were regenerated to
  match the new Gauss-Newton implementation.
- *Filmic Pro*: the look-up table domain start was changed from `0` to
  `EPSILON` and a `left=0` clamp was added.
- The `_SPOW_ENABLED` and `_SDIV_MODE` module states were migrated to
  `contextvars.ContextVar` for thread and async-task safety.
@KelSolaar KelSolaar force-pushed the feature/array-api-support branch from d825253 to dcdf17c Compare June 13, 2026 21:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant