Поддержка стандарта array API#

Примечание

Поддержка стандарта Array API всё ещё экспериментальна и скрыта за переменной окружения. На данный момент покрыта лишь небольшая часть публичного API.

Это руководство описывает, как использовать и добавить поддержку для the Стандарт Python array API. Этот стандарт позволяет пользователям использовать любую совместимую с Array API библиотеку массивов с частями SciPy из коробки.

The RFC определяет, как SciPy реализует поддержку стандарта, с основным принципом, заключающимся в «тип массива на входе равен типу массива на выходе». Кроме того, реализация выполняет более строгую проверку допустимых массивоподобных входных данных, например, отклоняя экземпляры numpy matrix и masked array, а также массивы с типом данных object.

В дальнейшем, совместимое с Array API пространство имен обозначается как xp.

Использование поддержки стандарта Array API#

Чтобы включить поддержку стандарта Array API, необходимо установить переменную окружения перед импортом SciPy:

export SCIPY_ARRAY_API=1

Это обеспечивает как поддержку стандарта Array API, так и более строгую проверку входных данных для аргументов типа массивов. Обратите внимание, что эта переменная окружения предназначена для временного использования, как способ внесения постепенных изменений и их слияния в ``main`` без немедленного влияния на обратную совместимость. Мы не намерены сохранять эту переменную окружения в долгосрочной перспективе.

Этот пример кластеризации показывает использование с тензорами PyTorch в качестве входных данных и возвращаемых значений:

>>> import torch
>>> from scipy.cluster.vq import vq
>>> code_book = torch.tensor([[1., 1., 1.],
...                           [2., 2., 2.]])
>>> features  = torch.tensor([[1.9, 2.3, 1.7],
...                           [1.5, 2.5, 2.2],
...                           [0.8, 0.6, 1.7]])
>>> code, dist = vq(features, code_book)
>>> code
tensor([1, 1, 0], dtype=torch.int32)
>>> dist
tensor([0.4359, 0.7348, 0.8307])

Обратите внимание, что приведенный выше пример работает для тензоров PyTorch CPU. Для тензоров GPU или массивов CuPy ожидаемый результат для vq является TypeError, потому что vq использует скомпилированный код в своей реализации, который не будет работать на GPU.

Более строгая проверка входных массивов будет отклонять np.matrix и np.ma.MaskedArray экземпляры, а также массивы с object dtype:

>>> import numpy as np
>>> from scipy.cluster.vq import vq
>>> code_book = np.array([[1., 1., 1.],
...                       [2., 2., 2.]])
>>> features  = np.array([[1.9, 2.3, 1.7],
...                       [1.5, 2.5, 2.2],
...                       [0.8, 0.6, 1.7]])
>>> vq(features, code_book)
(array([1, 1, 0], dtype=int32), array([0.43588989, 0.73484692, 0.83066239]))

>>> # The above uses numpy arrays; trying to use np.matrix instances or object
>>> # arrays instead will yield an exception with `SCIPY_ARRAY_API=1`:
>>> vq(np.asmatrix(features), code_book)
...
TypeError: 'numpy.matrix' are not supported

>>> vq(np.ma.asarray(features), code_book)
...
TypeError: 'numpy.ma.MaskedArray' are not supported

>>> vq(features.astype(np.object_), code_book)
...
TypeError: object arrays are not supported

В настоящее время поддерживаемая функциональность#

Следующие модули обеспечивают поддержку стандарта array API, когда установлена переменная окружения:

Отдельные функции в вышеуказанных модулях предоставляют таблицу возможностей в документации, подобную приведённой ниже. Если таблица отсутствует, функция ещё не поддерживает бэкенды, отличные от NumPy.

Таблица возможностей примеров#

Библиотека

CPU

GPU

NumPy

✅

н/д

CuPy

н/д

✅

PyTorch

✅

✅

JAX

⚠️ нет JIT

⛔

Dask

⛔

н/д

В приведённом выше примере функция имеет некоторую поддержку массивов NumPy, CuPy, PyTorch и JAX, но не поддерживает массивы Dask. Некоторые бэкенды, такие как JAX и PyTorch, изначально поддерживают несколько устройств (CPU и GPU), но поддержка SciPy таких массивов может быть ограничена; например, эта функция SciPy ожидается к работе только с массивами JAX, расположенными на CPU. Кроме того, некоторые бэкенды могут иметь серьёзные ограничения; в примере функция завершится ошибкой при запуске внутри jax.jit. Дополнительные предостережения могут быть перечислены в документации функции.

Хотя элементы таблицы, помеченные как «n/a», по своей природе выходят за рамки, мы постоянно работаем над заполнением остальных.

Пожалуйста, смотрите проблемный тикет для обновлений.

00:05.164#

Ключевая часть поддержки стандарта array API и специфических функций совместимости для Numpy, CuPy и PyTorch предоставляется через array-api-compat. Этот пакет включен в код SciPy через git submodule (под scipy/_lib), поэтому новые зависимости не вводятся.

array-api-compat предоставляет общие служебные функции и добавляет псевдонимы, такие как xp.concat (который, для numpy, сопоставлен с np.concatenate до того, как NumPy добавил np.concat в NumPy 2.0). Это позволяет использовать единый API в NumPy, PyTorch, CuPy и JAX (с другими библиотеками, такими как Dask, работа продолжается).

Когда переменная окружения не установлена и, следовательно, поддержка стандарта array API в SciPy отключена, мы все равно используем обернутую версию пространства имен NumPy, которая является array_api_compat.numpy. Это не должно изменять поведение функций SciPy, так как это, по сути, существующий numpy пространство имён с добавлением ряда псевдонимов и изменением/добавлением нескольких функций для поддержки стандарта Array API. Когда поддержка включена, xp = array_namespace(input) будет стандартно-совместимым пространством имён, соответствующим типу входного массива функции (например, если вход в cluster.vq.kmeans является тензором PyTorch, тогда xp является array_api_compat.torch).

Добавление поддержки стандарта Array API в функцию SciPy#

По возможности, новый код, добавляемый в SciPy, должен максимально следовать стандарту array API (эти функции обычно также являются лучшими практиками использования NumPy). Следуя стандарту, добавление поддержки стандарта array API обычно становится простым, и в идеале нам не нужно поддерживать никакие пользовательские настройки.

Различные вспомогательные функции доступны в scipy._lib._array_api - пожалуйста, смотрите __all__ в этом модуле для списка текущих вспомогательных функций и их документации для получения дополнительной информации.

Для добавления поддержки функции SciPy, которая определена в .py файл, что вам нужно изменить:

  1. Проверка входного массива,

  2. Используя xp скорее np функции,

  3. При вызове скомпилированного кода преобразуйте массив в массив NumPy до вызова и обратно в исходный тип массива после.

Проверка входного массива использует следующий шаблон:

xp = array_namespace(arr) # where arr is the input array
# alternatively, if there are multiple array inputs, include them all:
xp = array_namespace(arr1, arr2)

# replace np.asarray with xp.asarray
arr = xp.asarray(arr)
# uses of non-standard parameters of np.asarray can be replaced with _asarray
arr = _asarray(arr, order='C', dtype=xp.float64, xp=xp)

Обратите внимание, что если один вход является типом не-NumPy массива, все подобные массиву входы должны быть этого типа; попытка смешать не-NumPy массивы со списками, скалярами Python или другими произвольными объектами Python вызовет исключение. Для массивов NumPy эти типы будут продолжать приниматься по причинам обратной совместимости.

Если функция вызывает скомпилированный код только один раз, используйте следующий шаблон:

x = np.asarray(x)  # convert to numpy right before compiled call(s)
y = _call_compiled_code(x)
y = xp.asarray(y)  # convert back to original array type

Если есть несколько вызовов скомпилированного кода, убедитесь, что преобразование выполняется только один раз, чтобы избежать слишком больших накладных расходов.

Вот пример для гипотетической публичной функции SciPy toto:

def toto(a, b):
    a = np.asarray(a)
    b = np.asarray(b, copy=True)

    c = np.sum(a) - np.prod(b)

    # this is some C or Cython call
    d = cdist(c)

    return d

Вы бы преобразовали это так:

def toto(a, b):
    xp = array_namespace(a, b)
    a = xp.asarray(a)
    b = xp_copy(b, xp=xp)  # our custom helper is needed for copy

    c = xp.sum(a) - xp.prod(b)

    # this is some C or Cython call
    c = np.asarray(c)
    d = cdist(c)
    d = xp.asarray(d)

    return d

Проход через скомпилированный код требует возврата к массиву NumPy, потому что расширенные модули SciPy работают только с массивами NumPy (или memoryviews в случае Cython). Для массивов на CPU преобразования должны быть с нулевым копированием, в то время как на GPU и других устройствах попытка преобразования вызовет исключение. Причина в том, что скрытая передача данных между устройствами считается плохой практикой, так как это может быть большим и трудно обнаруживаемым узким местом производительности.

Добавление тестов#

Для запуска теста на нескольких бэкендах массивов, вам следует добавить xp фикстуру к нему, которая оценивается как текущее тестируемое пространство имен массивов.

Доступны следующие маркеры pytest:

  • skip_xp_backends(backend=None, reason=None, np_only=False, cpu_only=False, eager_only=False, exceptions=None): пропускать определённые бэкенды или категории бэкендов. См. строку документации scipy.conftest.skip_or_xfail_xp_backends для информации о том, как использовать этот маркер для пропуска тестов.

  • xfail_xp_backends(backend=None, reason=None, np_only=False, cpu_only=False, eager_only=False, exceptions=None): ожидать сбоя для определённых бэкендов или категорий бэкендов. См. документацию scipy.conftest.skip_or_xfail_xp_backends для информации о том, как использовать этот маркер для xfail тестов.

  • skip_xp_invalid_arg используется для пропуска тестов, которые используют аргументы, недействительные, когда SCIPY_ARRAY_API включён. Например, некоторые тесты scipy.stats функции передают маскированные массивы в тестируемую функцию, но маскированные массивы несовместимы с array API. Использование skip_xp_invalid_arg декоратор позволяет этим тестам защищаться от регрессий, когда SCIPY_ARRAY_API не используется без последствий, приводящих к сбоям, когда SCIPY_ARRAY_API используется. Со временем мы хотим, чтобы эти функции выдавали предупреждения об устаревании при получении недопустимых входных данных для array API, и этот декоратор проверит, что предупреждение об устаревании выдаётся без приведения к сбою теста. Когда SCIPY_ARRAY_API=1 поведение становится поведением по умолчанию и единственным поведением, эти тесты (и сам декоратор) будут удалены.

  • array_api_backends: этот маркер автоматически добавляется xp фикстуру для всех тестов, которые её используют. Это полезно, например, для выбора всех и только таких тестов:

    python dev.py test -b all -m array_api_backends

scipy._lib._array_api содержит утверждения, не зависящие от типа массива, такие как xp_assert_close который можно использовать для замены утверждений из numpy.testing.

Когда эти утверждения выполняются в тесте, который использует xp Фикстура обеспечивает соответствие пространств имён как фактического, так и ожидаемого массивов пространству имён, заданному фикстурой. Тесты без xp фикстура определяет пространство имен из желаемого массива. Этот механизм можно переопределить, явно передав xp= параметр для функций утверждения.

Следующие примеры демонстрируют, как использовать маркеры:

from scipy.conftest import skip_xp_invalid_arg
from scipy._lib._array_api import xp_assert_close
...
@pytest.mark.skip_xp_backends(np_only=True, reason='skip reason')
def test_toto1(self, xp):
    a = xp.asarray([1, 2, 3])
    b = xp.asarray([0, 2, 5])
    xp_assert_close(toto(a, b), a)
...
@pytest.mark.skip_xp_backends('array_api_strict', reason='skip reason 1')
@pytest.mark.skip_xp_backends('cupy', reason='skip reason 2')
def test_toto2(self, xp):
    ...
...
# Do not run when SCIPY_ARRAY_API is used
@skip_xp_invalid_arg
def test_toto_masked_array(self):
    ...

Передача имён бэкендов в exceptions означает, что они не будут пропущены функцией cpu_only=True или eager_only=True. Это полезно, когда делегирование реализовано для некоторых, но не всех, не-CPU бэкендов, а CPU-путь требует преобразования в NumPy для скомпилированного кода:

# array-api-strict and CuPy will always be skipped, for the given reasons.
# All libraries using a non-CPU device will also be skipped, apart from
# JAX, for which delegation is implemented (hence non-CPU execution is supported).
@pytest.mark.skip_xp_backends(cpu_only=True, exceptions=['jax.numpy'])
@pytest.mark.skip_xp_backends('array_api_strict', reason='skip reason 1')
@pytest.mark.skip_xp_backends('cupy', reason='skip reason 2')
def test_toto(self, xp):
    ...

После применения этих маркеров, dev.py test может использоваться с новой опцией -b или --array-api-backend:

python dev.py test -b numpy -b torch -s cluster

Это автоматически устанавливает SCIPY_ARRAY_API соответствующим образом. Чтобы протестировать библиотеку с несколькими устройствами с нестандартным устройством, вторая переменная окружения (SCIPY_DEVICE, используется только в тестовом наборе) может быть установлен. Допустимые значения зависят от используемой библиотеки массивов, например, для PyTorch допустимыми значениями являются "cpu", "cuda", "mps". Чтобы запустить набор тестов с бэкендом PyTorch MPS, используйте: SCIPY_DEVICE=mps python dev.py test -b torch.

Обратите внимание, что существует рабочий процесс GitHub Actions, который тестирует с array-api-strict, PyTorch и JAX на CPU.

Тестирование JAX JIT компилятора#

The JAX JIT компилятор вводит специальные ограничения для всего кода, обёрнутого @jax.jit, которые отсутствуют при запуске JAX в режиме eager. В частности, булевы маски в __getitem__ и в не поддерживаются, и вы не можете материализовать массивы, применяя bool(), float(), np.asarray() и т.д. к ним.

Для корректного тестирования scipy с JAX необходимо обернуть тестируемые функции scipy с @jax.jit до их вызова модульными тестами. Для этого следует пометить их следующим образом в вашем тестовом модуле:

from scipy._lib._lazy_testing import lazy_xp_function
from scipy.mymodule import toto

lazy_xp_function(toto)

def test_toto(xp):
    a = xp.asarray([1, 2, 3])
    b = xp.asarray([0, 2, 5])
    # When xp==jax.numpy, toto is wrapped with @jax.jit
    xp_assert_close(toto(a, b), a)

См. полную документацию в scipy/_lib/_lazy_testing.py.

Дополнительная информация#

Вот дополнительные ресурсы, которые мотивировали некоторые дизайнерские решения и помогли на этапе разработки:

  • Начальный PR с некоторыми обсуждениями

  • Быстрый старт с этого PR и некоторое вдохновение взято из scikit-learn.

  • PR добавление поддержки Array API в scikit-learn

  • Некоторые другие релевантные PR scikit-learn: #22554 и #25956