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.

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

Add generated doc for primitive and core ops to design page #351

Merged
merged 1 commit into from
Jan 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
6 changes: 5 additions & 1 deletion cubed/core/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,7 @@ def blockwise(


def elemwise(func, *args: "Array", dtype=None) -> "Array":
"""Apply a function elementwise to array arguments, respecting broadcasting."""
shapes = [arg.shape for arg in args]
out_ndim = len(np.broadcast_shapes(*shapes))
expr_inds = tuple(range(out_ndim))[::-1]
Expand All @@ -310,6 +311,7 @@ def elemwise(func, *args: "Array", dtype=None) -> "Array":


def index(x, key):
"Subset an array, along one or more axes."
if not isinstance(key, tuple):
key = (key,)

Expand Down Expand Up @@ -586,7 +588,7 @@ def map_direct(
func, *args: "Array", shape, dtype, chunks, extra_projected_mem, spec=None, **kwargs
) -> "Array":
"""
Map a function across blocks of a new array, using side-input arrays to read directly from.
Apply a function across blocks of a new array, reading directly from side inputs (not necessarily in a blockwise fashion).

Parameters
----------
Expand Down Expand Up @@ -736,6 +738,7 @@ def reduction(
keepdims=False,
extra_func_kwargs=None,
) -> "Array":
"""Apply a function to reduce an array along one or more axes."""
if combine_func is None:
combine_func = func
if axis is None:
Expand Down Expand Up @@ -826,6 +829,7 @@ def reduction(


def arg_reduction(x, /, arg_func, axis=None, *, keepdims=False):
"""A reduction that returns the array indexes, not the values."""
dtype = np.int64 # index data type
intermediate_dtype = [("i", dtype), ("v", x.dtype)]

Expand Down
2 changes: 1 addition & 1 deletion cubed/primitive/blockwise.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ def blockwise(
extra_func_kwargs: Optional[Dict[str, Any]] = None,
**kwargs,
):
"""Apply a function across blocks from multiple source Zarr arrays.
"""Apply a function to multiple blocks from multiple inputs, expressed using concise indexing rules.

Parameters
----------
Expand Down
2 changes: 1 addition & 1 deletion cubed/primitive/rechunk.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ def rechunk(
target_store: T_Store,
temp_store: Optional[T_Store] = None,
) -> List[CubedPipeline]:
"""Rechunk a Zarr array to have target_chunks.
"""Change the chunking of an array, without changing its shape or dtype.

Parameters
----------
Expand Down
44 changes: 24 additions & 20 deletions docs/design.md
Original file line number Diff line number Diff line change
Expand Up @@ -21,31 +21,35 @@ Cubed uses external runtimes for computation. It follows the Rechunker model (an

There are two primitive operations on arrays:

<dl>
<dt><code>blockwise</code></dt>
<dd>Applies a function to multiple blocks from multiple inputs, expressed using concise indexing rules.</dd>
<dt><code>rechunk</code></dt>
<dd>Changes the chunking of an array, without changing its shape or dtype.</dd>
</dl>
```{eval-rst}
.. currentmodule:: cubed.primitive.blockwise
.. autosummary::
:nosignatures:

blockwise
.. currentmodule:: cubed.primitive.rechunk
.. autosummary::
:nosignatures:

rechunk
```

## Core operations

These are built on top of the primitive operations, and provide functions that are needed to implement all array operations.

<dl>
<dt><code>elemwise</code></dt>
<dd>Applies a function elementwise to its arguments, respecting broadcasting.</dd>
<dt><code>map_blocks</code></dt>
<dd>Applies a function to corresponding blocks from multiple inputs.</dd>
<dt><code>map_direct</code></dt>
<dd>Applies a function across blocks of a new array, reading directly from side inputs (not necessarily in a blockwise fashion).</dd>
<dt><code>index</code> (<code>__getitem__</code>)</dt>
<dd>Subsets an array, along one or more axes.</dd>
<dt><code>reduction</code></dt>
<dd>Applies a function to reduce an array along one or more axes.</dd>
<dt><code>arg_reduction</code></dt>
<dd>A reduction that returns the array indexes, not the values.</dd>
</dl>
```{eval-rst}
.. currentmodule:: cubed.core.ops
.. autosummary::
:nosignatures:

elemwise
map_blocks
map_direct
index
reduction
arg_reduction
```

## Array API

Expand Down