Code

Add basic JAX support#84

Code

Merged

asmeurer merged 20 commits intodata-apis:maindata-apis/array-api-compat:mainfrom asmeurer:jaxasmeurer/array-api-compat:jaxCopy head branch name to clipboardFeb 9, 2024

+175-40

Commits

Commits on Feb 6, 2024

• Make is_numpy_array, is_cupy_array, is_torch_array, and is_dask_array public

  asmeurer

  committed
  2ee6902

  Copy full SHA for 2ee6902

• Note dask in a couple of places in the README

  asmeurer

  committed
  12b5294

  Copy full SHA for 12b5294

• Add JAX support

  Show description for 6bf5dad
  asmeurer

  committed
  6bf5dad

  Copy full SHA for 6bf5dad

• Install JAX on CI

  asmeurer

  committed
  583f6bb

  Copy full SHA for 583f6bb

• Install JAX as jax[cpu]

  asmeurer

  committed
  9c8bed6

  Copy full SHA for 9c8bed6

• Fix a comment

  asmeurer

  committed
  6d59ae8

  Copy full SHA for 6d59ae8

• Use jax.Array instead of jax.numpy.ndarray in is_jax_array

  asmeurer

  committed
  ce07cd9

  Copy full SHA for ce07cd9

• Use the native jax to_device() method

  asmeurer

  committed
  ddb313e

  Copy full SHA for ddb313e

• Add a basic test for device() and to_device()

  asmeurer

  committed
  6004b97

  Copy full SHA for 6004b97

• XFAIL the device test for dask.array

  asmeurer

  committed
  701a5ef

  Copy full SHA for 701a5ef

• Allow to_device(x, "cpu") for JAX arrays

  asmeurer

  committed
  049d557

  Copy full SHA for 049d557

• Move tests in test_common.py to test_helpers.py

  asmeurer

  committed
  db667ea

  Copy full SHA for db667ea

Commits on Feb 8, 2024

• Remove "cpu" device from jax to_device()

  asmeurer

  committed
  aafbbaa

  Copy full SHA for aafbbaa

• Merge branch 'main' into jax

  asmeurer

  committed
  919ec41

  Copy full SHA for 919ec41

• Rename import_or_skip_cupy to import_ and move jax logic into it

  asmeurer

  committed
  fa758f7

  Copy full SHA for fa758f7

• Skip JAX tests in Python 3.8

  asmeurer

  committed
  6c338ca

  Copy full SHA for 6c338ca

• Rename test_helpers.py to test_common.py

  Show description for 244462f
  asmeurer

  committed
  244462f

  Copy full SHA for 244462f

• Fix ruff warnings

  asmeurer

  committed
  bff9bf2

  Copy full SHA for bff9bf2

• Add jax.numpy to the test_array_namespace()

  Show description for 264e6c3
  asmeurer

  committed
  264e6c3

  Copy full SHA for 264e6c3

• Don't run the jax.numpy array_namespace test in Python 3.8

  asmeurer

  committed
  e7aff0f

  Copy full SHA for e7aff0f