Skip to content

Commit

Permalink
Change drop_axis to default to None rather than an empty list in …
Browse files Browse the repository at this point in the history
…`map_blocks`. (#446)

This follows Dask, see dask/dask#9462.
  • Loading branch information
tomwhite authored Apr 16, 2024
1 parent 1935e99 commit 02fa89f
Showing 1 changed file with 14 additions and 2 deletions.
16 changes: 14 additions & 2 deletions cubed/core/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -564,13 +564,16 @@ def map_blocks(
*args: "Array",
dtype=None,
chunks=None,
drop_axis=[],
drop_axis=None,
new_axis=None,
spec=None,
**kwargs,
) -> "Array":
"""Apply a function to corresponding blocks from multiple input arrays."""

if drop_axis is None:
drop_axis = []

# Handle the case where an array is created by calling `map_blocks` with no input arrays
if len(args) == 0:
from cubed.array_api.creation_functions import empty_virtual_array
Expand Down Expand Up @@ -618,10 +621,19 @@ def wrap(*a, **kw):


def _map_blocks(
func, *args: "Array", dtype=None, chunks=None, drop_axis=[], new_axis=None, **kwargs
func,
*args: "Array",
dtype=None,
chunks=None,
drop_axis=None,
new_axis=None,
**kwargs,
) -> "Array":
# based on dask

if drop_axis is None:
drop_axis = []

new_axes = {}

if isinstance(drop_axis, Number):
Expand Down

0 comments on commit 02fa89f

Please sign in to comment.