From fae50aa0b67d806ed3b193e904429a988b4b1cdd Mon Sep 17 00:00:00 2001 From: Illviljan <14371165+Illviljan@users.noreply.github.com> Date: Sun, 22 Sep 2024 00:38:38 +0200 Subject: [PATCH] add examples --- .../_array_api/_linear_algebra_functions.py | 27 +++++++++++++++++++ xarray/namedarray/_array_api/_utils.py | 16 ++++++----- 2 files changed, 36 insertions(+), 7 deletions(-) diff --git a/xarray/namedarray/_array_api/_linear_algebra_functions.py b/xarray/namedarray/_array_api/_linear_algebra_functions.py index f8b730907af..f5e2cc0f389 100644 --- a/xarray/namedarray/_array_api/_linear_algebra_functions.py +++ b/xarray/namedarray/_array_api/_linear_algebra_functions.py @@ -10,6 +10,33 @@ def matmul( x1: NamedArray[Any, Any], x2: NamedArray[Any, Any], / ) -> NamedArray[Any, Any]: + """ + Matrix product of two arrays. + + Examples + -------- + For 2-D arrays it is the matrix product: + + >>> import numpy as np + >>> a = NamedArray(("y", "x"), np.array([[1, 0], [0, 1]])) + >>> b = NamedArray(("y", "x"), np.array([[4, 1], [2, 2]])) + >>> matmul(a, b) + + + For 2-D mixed with 1-D, the result is the usual. + + >>> a = NamedArray(("y", "x"), np.array([[1, 0], [0, 1]])) + >>> b = NamedArray(("x",), np.array([1, 2])) + >>> matmul(a, b) + + Broadcasting is conventional for stacks of arrays + + >>> a = NamedArray(("z", "y", "x"), np.arange(2 * 2 * 4).reshape((2, 2, 4))) + >>> b = NamedArray(("z", "y", "x"), np.arange(2 * 2 * 4).reshape((2, 4, 2))) + >>> axb = matmul(a,b) + >>> axb.dims, axb.shape + """ xp = _get_data_namespace(x1) _data = xp.matmul(x1._data, x2._data) # TODO: Figure out a better way: diff --git a/xarray/namedarray/_array_api/_utils.py b/xarray/namedarray/_array_api/_utils.py index 30a125d5f6c..c0a6611427d 100644 --- a/xarray/namedarray/_array_api/_utils.py +++ b/xarray/namedarray/_array_api/_utils.py @@ -111,6 +111,8 @@ def _infer_dims( >>> _infer_dims((3, 1)) ('dim_1', 'dim_0') + >>> _infer_dims((), ()) + () >>> _infer_dims((1,), "x") ('x',) >>> _infer_dims((1,), None) @@ -406,19 +408,19 @@ def _dims_from_tuple_indexing(dims: _Dims, key: _IndexKeys) -> _Dims: Examples -------- - >>> dims_from_tuple_indexing(("x", "y"), ()) + >>> _dims_from_tuple_indexing(("x", "y"), ()) ('x', 'y') - >>> dims_from_tuple_indexing(("x", "y"), (0,)) + >>> _dims_from_tuple_indexing(("x", "y"), (0,)) ('y',) - >>> dims_from_tuple_indexing(("x", "y"), (0, 0)) + >>> _dims_from_tuple_indexing(("x", "y"), (0, 0)) () - >>> dims_from_tuple_indexing(("x", "y"), (0, ...)) + >>> _dims_from_tuple_indexing(("x", "y"), (0, ...)) ('y',) - >>> dims_from_tuple_indexing(("x", "y"), (0, slice(0))) + >>> _dims_from_tuple_indexing(("x", "y"), (0, slice(0))) ('y',) - >>> dims_from_tuple_indexing(("x", "y"), (None,)) + >>> _dims_from_tuple_indexing(("x", "y"), (None,)) ('dim_2', 'x', 'y') - >>> dims_from_tuple_indexing(("x", "y"), (0, None, None, 0)) + >>> _dims_from_tuple_indexing(("x", "y"), (0, None, None, 0)) ('dim_1', 'dim_2') """ _dims = list(dims)