Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Expose the linalg namespace and include in status page #581

Merged
merged 6 commits into from
Sep 24, 2024
Merged

Expose the linalg namespace and include in status page #581

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e82a0df
Add matmul, matrix_transpose, tensordot, vecdot to linalg namespace
tomwhite Sep 22, 2024
0846e31
Move outer to linalg namespace
tomwhite Sep 22, 2024
091d41e
Remove flip from list of unimplemented functions since it was added i…
tomwhite Sep 22, 2024
3e5ad0a
Remove unstack from list of unimplemented functions since it was adde…
tomwhite Sep 22, 2024
1e5ff48
Add link to cumulative_sum PR
tomwhite Sep 22, 2024
2ad2f97
Add linalg table to status page
tomwhite Sep 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 33 additions & 3 deletions api_status.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
## Array API Coverage Implementation Status

Cubed supports version [2022.12](https://data-apis.org/array-api/2022.12/index.html) of the Python array API standard, with a few exceptions noted below. The [linear algebra extensions](https://data-apis.org/array-api/2022.12/extensions/linear_algebra_functions.html) and [Fourier transform functions](https://data-apis.org/array-api/2022.12/extensions/fourier_transform_functions.html) are *not* supported.
Cubed supports version [2022.12](https://data-apis.org/array-api/2022.12/index.html) of the Python array API standard, with a few exceptions noted below. The [Fourier transform functions](https://data-apis.org/array-api/2022.12/extensions/fourier_transform_functions.html) are *not* supported.

Support for version [2023.12](https://data-apis.org/array-api/2023.12/index.html) is tracked in Cubed issue [#438](/~https://github.com/cubed-dev/cubed/issues/438).

Expand Down Expand Up @@ -67,7 +67,7 @@ This table shows which parts of the the [Array API](https://data-apis.org/array-
| | `squeeze` | :white_check_mark: | | |
| | `stack` | :white_check_mark: | | |
| | `tile` | :x: | 2023.12 | |
| | `unstack` | :x: | 2023.12 | |
| | `unstack` | :white_check_mark: | 2023.12 | |
| Searching Functions | `argmax` | :white_check_mark: | | |
| | `argmin` | :white_check_mark: | | |
| | `nonzero` | :x: | | Shape is data dependent |
Expand All @@ -79,7 +79,7 @@ This table shows which parts of the the [Array API](https://data-apis.org/array-
| | `unique_values` | :x: | | Shape is data dependent |
| Sorting Functions | `argsort` | :x: | | Not in Dask |
| | `sort` | :x: | | Not in Dask |
| Statistical Functions | `cumulative_sum` | :x: | 2023.12 | |
| Statistical Functions | `cumulative_sum` | :x: | 2023.12 | WIP [#531](/~https://github.com/cubed-dev/cubed/pull/531) |
| | `max` | :white_check_mark: | | |
| | `mean` | :white_check_mark: | | |
| | `min` | :white_check_mark: | | |
Expand All @@ -89,3 +89,33 @@ This table shows which parts of the the [Array API](https://data-apis.org/array-
| | `var` | :x: | | Like `mean`, [#29](/~https://github.com/cubed-dev/cubed/issues/29) |
| Utility Functions | `all` | :white_check_mark: | | |
| | `any` | :white_check_mark: | | |

### Linear Algebra Extension

A few of the [linear algebra extension](https://data-apis.org/array-api/2022.12/extensions/linear_algebra_functions.html) functions are supported, as indicated in this table.

| Category | Object/Function | Implemented | Version | Notes |
| ------------------------ | ------------------- | ------------------ | ---------- | ---------------------------- |
| Linear Algebra Functions | `cholesky` | :x: | | |
| | `cross` | :x: | | |
| | `det` | :x: | | |
| | `diagonal` | :x: | | |
| | `eigh` | :x: | | |
| | `eigvalsh` | :x: | | |
| | `inv` | :x: | | |
| | `matmul` | :white_check_mark: | | |
| | `matrix_norm` | :x: | | |
| | `matrix_power` | :x: | | |
| | `matrix_rank` | :x: | | |
| | `matrix_transpose` | :white_check_mark: | | |
| | `outer` | :white_check_mark: | | |
| | `pinv` | :x: | | |
| | `qr` | :white_check_mark: | | |
| | `slogdet` | :x: | | |
| | `solve` | :x: | | |
| | `svd` | :x: | | |
| | `svdvals` | :x: | | |
| | `tensordot` | :white_check_mark: | | |
| | `trace` | :x: | | |
| | `vecdot` | :white_check_mark: | | |
| | `vectornorm` | :x: | | |
3 changes: 1 addition & 2 deletions cubed/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -267,12 +267,11 @@
from .array_api.linear_algebra_functions import (
matmul,
matrix_transpose,
outer,
tensordot,
vecdot,
)

__all__ += ["matmul", "matrix_transpose", "outer", "tensordot", "vecdot"]
__all__ += ["matmul", "matrix_transpose", "tensordot", "vecdot"]

from .array_api.manipulation_functions import (
broadcast_arrays,
Expand Down
4 changes: 2 additions & 2 deletions cubed/array_api/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,9 +212,9 @@

__all__ += ["take"]

from .linear_algebra_functions import matmul, matrix_transpose, outer, tensordot, vecdot
from .linear_algebra_functions import matmul, matrix_transpose, tensordot, vecdot

__all__ += ["matmul", "matrix_transpose", "outer", "tensordot", "vecdot"]
__all__ += ["matmul", "matrix_transpose", "tensordot", "vecdot"]

from .manipulation_functions import (
broadcast_arrays,
Expand Down
14 changes: 13 additions & 1 deletion cubed/array_api/linalg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,23 @@
from typing import NamedTuple

from cubed.array_api.array_object import Array

# These functions are in both the main and linalg namespaces
from cubed.array_api.linear_algebra_functions import ( # noqa: F401
matmul,
matrix_transpose,
tensordot,
vecdot,
)
from cubed.backend_array_api import namespace as nxp
from cubed.core.ops import general_blockwise, map_direct, merge_chunks
from cubed.core.ops import blockwise, general_blockwise, map_direct, merge_chunks
from cubed.utils import array_memory, get_item


def outer(x1, x2, /):
return blockwise(nxp.linalg.outer, "ij", x1, "i", x2, "j", dtype=x1.dtype)


class QRResult(NamedTuple):
Q: Array
R: Array
Expand Down
4 changes: 0 additions & 4 deletions cubed/array_api/linear_algebra_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,10 +95,6 @@ def matrix_transpose(x, /):
return permute_dims(x, axes)


def outer(x1, x2, /):
return blockwise(nxp.linalg.outer, "ij", x1, "i", x2, "j", dtype=x1.dtype)


def tensordot(x1, x2, /, *, axes=2, use_new_impl=True, split_every=None):
from cubed.array_api.statistical_functions import sum

Expand Down
2 changes: 1 addition & 1 deletion cubed/tests/test_array_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -417,7 +417,7 @@ def test_matmul_modal(modal_executor):
def test_outer(spec, executor):
a = xp.asarray([0, 1, 2], chunks=2, spec=spec)
b = xp.asarray([10, 50, 100], chunks=2, spec=spec)
c = xp.outer(a, b)
c = xp.linalg.outer(a, b)
assert_array_equal(c.compute(executor=executor), np.outer([0, 1, 2], [10, 50, 100]))


Expand Down
1 change: 0 additions & 1 deletion docs/array-api.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ The following parts of the standard are not implemented:
| Array object | In-place Ops |
| Creation Functions | `from_dlpack` |
| Indexing | Boolean array |
| Manipulation Functions | `flip` |
| Searching Functions | `nonzero` |
| Set Functions | `unique_all` |
| | `unique_counts` |
Expand Down
Loading