Add TensorFlow array API backend

Author: njzjz

.github/workflows/array-api-tests-tensorflow.yml

@@ -0,0 +1,13 @@
+name: Array API Tests (TensorFlow)
+
+on: [push, pull_request]
+
+jobs:
+  array-api-tests-tensorflow:
+    uses: ./.github/workflows/array-api-tests.yml
+    with:
+      package-name: tensorflow
+      python-versions: '[''3.10'', ''3.13'']'
+      pytest-extra-args: -n 2
+      extra-env-vars: |
+        ARRAY_API_TESTS_XFAIL_MARK=skip

meson.build

@@ -61,6 +61,15 @@ sources_raw = {
     'src/array_api_compat/torch/fft.py',
     'src/array_api_compat/torch/linalg.py',
   ],
+
+  'array_api_compat/tensorflow': [
+    'src/array_api_compat/tensorflow/__init__.py',
+    'src/array_api_compat/tensorflow/_aliases.py',
+    'src/array_api_compat/tensorflow/_info.py',
+    'src/array_api_compat/tensorflow/_typing.py',
+    'src/array_api_compat/tensorflow/fft.py',
+    'src/array_api_compat/tensorflow/linalg.py',
+  ],
 }
 
 sources = {}

pyproject.toml

@@ -38,6 +38,7 @@ numpy = ["numpy>=1.22"]
 pytorch = ["torch"]
 sparse = ["sparse>=0.15.1"]
 ndonnx = ["ndonnx"]
+tensorflow = ["tensorflow"]
 docs = [
     "furo",
     "linkify-it-py",
@@ -55,6 +56,7 @@ dev = [
     "pytest",
     "torch",
     "sparse>=0.15.1",
+    "tensorflow",
 ]
 
 
@@ -105,7 +107,7 @@ warn_unused_ignores = true
 warn_unreachable = true
 
 [[tool.mypy.overrides]]
-module = ["cupy.*", "cupy_backends.*", "dask.*", "jax.*", "ndonnx.*", "sparse.*", "torch.*"]
+module = ["cupy.*", "cupy_backends.*", "dask.*", "jax.*", "ndonnx.*", "sparse.*", "tensorflow.*", "torch.*"]
 ignore_missing_imports = true
 
 

src/array_api_compat/common/_helpers.py

@@ -37,6 +37,7 @@
     import numpy as np
     import numpy.typing as npt
     import sparse
+    import tensorflow as tf
     import torch
 
     # TODO: import from typing (requires Python >=3.13)
@@ -51,6 +52,7 @@
         | jax.Array
         | ndx.Array
         | sparse.SparseArray
+        | tf.Tensor
         | torch.Tensor
         | SupportsArrayNamespace[Any]
     )
@@ -168,6 +170,29 @@ def is_torch_array(x: object) -> TypeIs[torch.Tensor]:
     return _issubclass_fast(cls, "torch", "Tensor")
 
 
+def is_tensorflow_array(x: object) -> TypeIs[tf.Tensor]:
+    """
+    Return True if `x` is a TensorFlow tensor.
+
+    This function does not import TensorFlow if it has not already been
+    imported and is therefore cheap to use.
+
+    See Also
+    --------
+
+    array_namespace
+    is_array_api_obj
+    is_numpy_array
+    is_cupy_array
+    is_torch_array
+    is_dask_array
+    is_jax_array
+    is_pydata_sparse_array
+    """
+    cls = cast(Hashable, type(x))
+    return _issubclass_fast(cls, "tensorflow", "Tensor")
+
+
 def is_ndonnx_array(x: object) -> TypeIs[ndx.Array]:
     """
     Return True if `x` is a ndonnx Array.
@@ -308,6 +333,7 @@ def _is_array_api_cls(cls: type) -> bool:
         or _issubclass_fast(cls, "numpy", "generic")
         or _issubclass_fast(cls, "cupy", "ndarray")
         or _issubclass_fast(cls, "torch", "Tensor")
+        or _issubclass_fast(cls, "tensorflow", "Tensor")
         or _issubclass_fast(cls, "dask.array", "Array")
         or _issubclass_fast(cls, "sparse", "SparseArray")
         # TODO: drop support for jax<0.4.32 which didn't have __array_namespace__
@@ -387,6 +413,30 @@ def is_torch_namespace(xp: Namespace) -> bool:
     return xp.__name__ in {"torch", _compat_module_name() + ".torch"}
 
 
+@lru_cache(100)
+def is_tensorflow_namespace(xp: Namespace) -> bool:
+    """
+    Returns True if `xp` is a TensorFlow namespace.
+
+    This includes both TensorFlow itself and the version wrapped by
+    array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {"tensorflow", _compat_module_name() + ".tensorflow"}
+
+
 def is_ndonnx_namespace(xp: Namespace) -> bool:
     """
     Returns True if `xp` is an NDONNX namespace.
@@ -551,6 +601,14 @@ def _cls_to_namespace(
             import torch as xp  # type: ignore[no-redef]
         return xp, None
 
+    if _issubclass_fast(cls_, "tensorflow", "Tensor"):
+        if _use_compat:
+            _check_api_version(api_version)
+            from .. import tensorflow as xp  # type: ignore[no-redef]
+        else:
+            import tensorflow as xp  # type: ignore[no-redef]
+        return xp, None
+
     if _issubclass_fast(cls_, "dask.array", "Array"):
         if _use_compat:
             _check_api_version(api_version)
@@ -790,6 +848,8 @@ def device(x: _ArrayApiObj, /) -> Device:
             return x_device()
         else:
             return x_device
+    elif is_tensorflow_array(x):
+        return x.device
     elif is_pydata_sparse_array(x):
         # `sparse` will gain `.device`, so check for this first.
         x_device = getattr(x, "device", None)
@@ -851,6 +911,20 @@ def _torch_to_device(
     return x.to(device)
 
 
+def _tensorflow_to_device(
+    x: tf.Tensor,
+    device: str,
+    /,
+    stream: int | Any | None = None,
+) -> tf.Tensor:
+    if stream is not None:
+        raise ValueError("The stream argument to to_device() is not supported")
+    import tensorflow as tf
+
+    with tf.device(device):
+        return tf.identity(x)
+
+
 def to_device(x: Array, device: Device, /, *, stream: int | Any | None = None) -> Array:
     """
     Copy the array from the device on which it currently resides to the specified ``device``.
@@ -911,6 +985,8 @@ def to_device(x: Array, device: Device, /, *, stream: int | Any | None = None) -
         return _cupy_to_device(x, device, stream=stream)
     elif is_torch_array(x):
         return _torch_to_device(x, device, stream=stream)
+    elif is_tensorflow_array(x):
+        return _tensorflow_to_device(x, device, stream=stream)
     elif is_dask_array(x):
         if stream is not None:
             raise ValueError("The stream argument to to_device() is not supported")
@@ -965,6 +1041,7 @@ def _is_writeable_cls(cls: type) -> bool | None:
         or _issubclass_fast(cls, "jax", "Array")
         or _issubclass_fast(cls, "jax.core", "Tracer")  # see is_jax_array for limitations
         or _issubclass_fast(cls, "sparse", "SparseArray")
+        or _issubclass_fast(cls, "tensorflow", "Tensor")
     ):
         return False
     if _is_array_api_cls(cls):
@@ -998,6 +1075,7 @@ def _is_lazy_cls(cls: type) -> bool | None:
         or _issubclass_fast(cls, "numpy", "generic")
         or _issubclass_fast(cls, "cupy", "ndarray")
         or _issubclass_fast(cls, "torch", "Tensor")
+        or _issubclass_fast(cls, "tensorflow", "Tensor")
         or _issubclass_fast(cls, "sparse", "SparseArray")
     ):
         return False
@@ -1083,6 +1161,8 @@ def is_lazy_array(x: object) -> TypeGuard[_ArrayApiObj]:
     "is_numpy_namespace",
     "is_torch_array",
     "is_torch_namespace",
+    "is_tensorflow_array",
+    "is_tensorflow_namespace",
     "is_ndonnx_array",
     "is_ndonnx_namespace",
     "is_pydata_sparse_array",