dask-array 0.1.0__py3-none-any.whl

dask_array/_dispatch.py

@@ -0,0 +1,141 @@
+"""
+Dispatch registries for dask_array.
+
+This module provides Dispatch objects for array operations that need to be
+dispatched based on array type (numpy, cupy, sparse, etc.).
+
+concatenate_lookup and tensordot_lookup are defined in _core_utils.py but
+re-exported here for convenience.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from dask.utils import Dispatch
+
+# Re-export from _core_utils for convenience
+from dask_array._core_utils import concatenate_lookup, tensordot_lookup
+
+# Dispatch registries for array operations
+take_lookup = Dispatch("take")
+einsum_lookup = Dispatch("einsum")
+empty_lookup = Dispatch("empty")
+divide_lookup = Dispatch("divide")
+percentile_lookup = Dispatch("percentile")
+numel_lookup = Dispatch("numel")
+nannumel_lookup = Dispatch("nannumel")
+
+
+# --- numpy implementations ---
+
+
+def _divide(x1, x2, out=None, dtype=None):
+    """Implementation of numpy.divide that works with dtype kwarg."""
+    x = np.divide(x1, x2, out)
+    if dtype is not None:
+        x = x.astype(dtype)
+    return x
+
+
+def _percentile(a, q, method="linear"):
+    """
+    Chunk-level percentile calculation.
+
+    Returns (percentile_values, n) tuple where n is the number of elements.
+    Used for combining percentiles from multiple chunks.
+    """
+    from collections.abc import Iterator
+
+    n = len(a)
+    if not len(a):
+        return None, n
+    if isinstance(q, Iterator):
+        q = list(q)
+    if a.dtype.name == "category":
+        result = np.percentile(a.cat.codes, q, method=method)
+        import pandas as pd
+
+        return (
+            pd.Categorical.from_codes(result, a.dtype.categories, a.dtype.ordered),
+            n,
+        )
+    if type(a.dtype).__name__ == "DatetimeTZDtype":
+        import pandas as pd
+
+        if isinstance(a, (pd.Series, pd.Index)):
+            a = a.values
+
+    if np.issubdtype(a.dtype, np.datetime64):
+        values = a
+        if type(a).__name__ in ("Series", "Index"):
+            a2 = values.astype("i8")
+        else:
+            a2 = values.view("i8")
+        result = np.percentile(a2, q, method=method).astype(values.dtype)
+        if q[0] == 0:
+            # https://github.com/dask/dask/issues/6864
+            result[0] = min(result[0], values.min())
+        return result, n
+    if not np.issubdtype(a.dtype, np.number):
+        method = "nearest"
+    return np.percentile(a, q, method=method), n
+
+
+def _numel(x, **kwargs):
+    """
+    A reduction to count the number of elements.
+
+    Returns ndarray result (coerces to numpy).
+    """
+    import math
+
+    shape = x.shape
+    keepdims = kwargs.get("keepdims", False)
+    axis = kwargs.get("axis")
+    dtype = kwargs.get("dtype", np.float64)
+
+    if axis is None:
+        prod = np.prod(shape, dtype=dtype)
+        if keepdims is False:
+            return prod
+
+        return np.full(shape=(1,) * len(shape), fill_value=prod, dtype=dtype)
+
+    if not isinstance(axis, (tuple, list)):
+        axis = [axis]
+
+    prod = math.prod(shape[dim] for dim in axis)
+    if keepdims is True:
+        new_shape = tuple(shape[dim] if dim not in axis else 1 for dim in range(len(shape)))
+    else:
+        new_shape = tuple(shape[dim] for dim in range(len(shape)) if dim not in axis)
+
+    return np.broadcast_to(np.array(prod, dtype=dtype), new_shape)
+
+
+def _nannumel(x, **kwargs):
+    """A reduction to count the number of elements, excluding nans"""
+    return np.sum(~(np.isnan(x)), **kwargs)
+
+
+# --- Register numpy implementations ---
+
+take_lookup.register((object, np.ndarray, np.ma.masked_array), np.take)
+einsum_lookup.register((object, np.ndarray), np.einsum)
+empty_lookup.register((object, np.ndarray), np.empty)
+empty_lookup.register(np.ma.masked_array, np.ma.empty)
+divide_lookup.register((object, np.ndarray), _divide)
+divide_lookup.register(np.ma.masked_array, np.ma.divide)
+percentile_lookup.register(np.ndarray, _percentile)
+numel_lookup.register((object, np.ndarray), _numel)
+nannumel_lookup.register((object, np.ndarray), _nannumel)
+
+
+# --- Register masked array numel ---
+
+
+@numel_lookup.register(np.ma.masked_array)
+def _numel_masked(x, **kwargs):
+    """Numel implementation for masked arrays."""
+    return np.sum(np.ones_like(x), **kwargs)

dask_array/_einsum.py

@@ -0,0 +1,277 @@
+"""Einstein summation for array-expr."""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+from dask import config
+from dask.utils import cached_max, derived_from
+
+from dask_array._dispatch import einsum_lookup
+
+# Valid characters for einsum subscripts (from numpy)
+einsum_symbols = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+einsum_symbols_set = set(einsum_symbols)
+
+
+def chunk_einsum(*operands, **kwargs):
+    """Chunk-level einsum computation.
+
+    This function is used by blockwise to compute einsum on individual chunks.
+    It dispatches to the appropriate einsum implementation based on array type.
+    """
+    subscripts = kwargs.pop("subscripts")
+    ncontract_inds = kwargs.pop("ncontract_inds")
+    dtype = kwargs.pop("kernel_dtype")
+    einsum = einsum_lookup.dispatch(type(operands[0]))
+    chunk = einsum(subscripts, *operands, dtype=dtype, **kwargs)
+
+    # Avoid concatenate=True in blockwise by adding 1's
+    # for the contracted dimensions
+    return chunk.reshape(chunk.shape + (1,) * ncontract_inds)
+
+
+def _calculate_new_chunksizes(old_chunks, new_chunks, changeable_dimensions, target_size):
+    """Calculate new chunk sizes for einsum rechunking."""
+    from dask_array._shuffle import _calculate_new_chunksizes as _calc
+
+    return _calc(old_chunks, new_chunks, changeable_dimensions, target_size)
+
+
+def _parse_einsum_input(operands, asarray):
+    """Parse einsum input, adapted from numpy/dask.
+
+    This is a copy of parse_einsum_input from einsumfuncs.py but uses
+    the provided asarray function to ensure correct array type.
+    """
+    if len(operands) == 0:
+        raise ValueError("No input operands")
+
+    if isinstance(operands[0], str):
+        subscripts = operands[0].replace(" ", "")
+        operands = [asarray(o) for o in operands[1:]]
+
+        # Ensure all characters are valid
+        for s in subscripts:
+            if s in ".,->":
+                continue
+            if s not in einsum_symbols_set:
+                raise ValueError(f"Character {s} is not a valid symbol.")
+
+    else:
+        tmp_operands = list(operands)
+        operand_list = []
+        subscript_list = []
+        for _ in range(len(operands) // 2):
+            operand_list.append(tmp_operands.pop(0))
+            subscript_list.append(tmp_operands.pop(0))
+
+        output_list = tmp_operands[-1] if len(tmp_operands) else None
+        operands = [asarray(v) for v in operand_list]
+        subscripts = ""
+        last = len(subscript_list) - 1
+        for num, sub in enumerate(subscript_list):
+            for s in sub:
+                if s is Ellipsis:
+                    subscripts += "..."
+                elif isinstance(s, int):
+                    subscripts += einsum_symbols[s]
+                else:
+                    raise TypeError("For this input type lists must contain either int or Ellipsis")
+            if num != last:
+                subscripts += ","
+
+        if output_list is not None:
+            subscripts += "->"
+            for s in output_list:
+                if s is Ellipsis:
+                    subscripts += "..."
+                elif isinstance(s, int):
+                    subscripts += einsum_symbols[s]
+                else:
+                    raise TypeError("For this input type lists must contain either int or Ellipsis")
+    # Check for proper "->"
+    if ("-" in subscripts) or (">" in subscripts):
+        invalid = (subscripts.count("-") > 1) or (subscripts.count(">") > 1)
+        if invalid or (subscripts.count("->") != 1):
+            raise ValueError("Subscripts can only contain one '->'.")
+
+    # Parse ellipses
+    if "." in subscripts:
+        used = subscripts.replace(".", "").replace(",", "").replace("->", "")
+        unused = list(einsum_symbols_set - set(used))
+        ellipse_inds = "".join(unused)
+        longest = 0
+
+        if "->" in subscripts:
+            input_tmp, output_sub = subscripts.split("->")
+            split_subscripts = input_tmp.split(",")
+            out_sub = True
+        else:
+            split_subscripts = subscripts.split(",")
+            out_sub = False
+
+        for num, sub in enumerate(split_subscripts):
+            if "." in sub:
+                if (sub.count(".") != 3) or (sub.count("...") != 1):
+                    raise ValueError("Invalid Ellipses.")
+
+                # Take into account numerical values
+                if operands[num].shape == ():
+                    ellipse_count = 0
+                else:
+                    ellipse_count = max(operands[num].ndim, 1)
+                    ellipse_count -= len(sub) - 3
+
+                if ellipse_count > longest:
+                    longest = ellipse_count
+
+                if ellipse_count < 0:
+                    raise ValueError("Ellipses lengths do not match.")
+                elif ellipse_count == 0:
+                    split_subscripts[num] = sub.replace("...", "")
+                else:
+                    rep_inds = ellipse_inds[-ellipse_count:]
+                    split_subscripts[num] = sub.replace("...", rep_inds)
+
+        subscripts = ",".join(split_subscripts)
+        if longest == 0:
+            out_ellipse = ""
+        else:
+            out_ellipse = ellipse_inds[-longest:]
+
+        if out_sub:
+            subscripts += "->" + output_sub.replace("...", out_ellipse)
+        else:
+            # Special care for outputless ellipses
+            output_subscript = ""
+            tmp_subscripts = subscripts.replace(",", "")
+            for s in sorted(set(tmp_subscripts)):
+                if s not in einsum_symbols_set:
+                    raise ValueError(f"Character {s} is not a valid symbol.")
+                if tmp_subscripts.count(s) == 1:
+                    output_subscript += s
+            normal_inds = "".join(sorted(set(output_subscript) - set(out_ellipse)))
+
+            subscripts += f"->{out_ellipse}{normal_inds}"
+
+    # Build output string if does not exist
+    if "->" in subscripts:
+        input_subscripts, output_subscript = subscripts.split("->")
+    else:
+        input_subscripts = subscripts
+        # Build output subscripts
+        tmp_subscripts = subscripts.replace(",", "")
+        output_subscript = ""
+        for s in sorted(set(tmp_subscripts)):
+            if s not in einsum_symbols_set:
+                raise ValueError(f"Character {s} is not a valid symbol.")
+            if tmp_subscripts.count(s) == 1:
+                output_subscript += s
+
+    # Make sure output subscripts are in the input
+    for char in output_subscript:
+        if char not in input_subscripts:
+            raise ValueError(f"Output character {char} did not appear in the input")
+
+    # Make sure number operands is equivalent to the number of terms
+    if len(input_subscripts.split(",")) != len(operands):
+        raise ValueError("Number of einsum subscripts must be equal to the number of operands.")
+
+    return (input_subscripts, output_subscript, operands)
+
+
+@derived_from(np)
+def einsum(*operands, dtype=None, optimize=False, split_every=None, **kwargs):
+    """Dask added an additional keyword-only argument ``split_every``.
+
+    split_every: int >= 2 or dict(axis: int), optional
+        Determines the depth of the recursive aggregation.
+        Defaults to ``None`` which would let dask heuristically
+        decide a good default.
+    """
+    from dask_array._collection import asarray, blockwise
+
+    einsum_dtype = dtype
+
+    # Parse operands, converting to dask arrays using array-expr asarray
+    inputs, outputs, ops = _parse_einsum_input(operands, asarray)
+
+    subscripts = "->".join((inputs, outputs))
+
+    # Infer the output dtype from operands
+    if dtype is None:
+        dtype = np.result_type(*[o.dtype for o in ops])
+
+    if optimize is not False:
+        # Avoid computation of dask arrays within np.einsum_path
+        # by passing in small numpy arrays broadcasted
+        # up to the right shape
+        fake_ops = [np.broadcast_to(o.dtype.type(0), shape=o.shape) for o in ops]
+        optimize, _ = np.einsum_path(subscripts, *fake_ops, optimize=optimize)
+
+    inputs = [tuple(i) for i in inputs.split(",")]
+
+    # Set of all indices
+    all_inds = {a for i in inputs for a in i}
+
+    # Which indices are contracted?
+    contract_inds = all_inds - set(outputs)
+    ncontract_inds = len(contract_inds)
+
+    if len(inputs) > 1 and len(outputs) > 0:
+        # Calculate the increase in chunk size compared to the largest input chunk
+        max_chunk_sizes, max_chunk_size_input = {}, 1
+        for op, input in zip(ops, inputs):
+            max_chunk_size_input = max(math.prod(map(cached_max, op.chunks)), max_chunk_size_input)
+            max_chunk_sizes.update(
+                {
+                    inp: max(cached_max(op.chunks[i]), max_chunk_sizes.get(inp, 1))
+                    for i, inp in enumerate(input)
+                    if inp not in contract_inds
+                }
+            )
+
+        max_chunk_size_output = math.prod(max_chunk_sizes.values())
+        factor = max_chunk_size_output / (max_chunk_size_input * config.get("array.chunk-size-tolerance"))
+
+        # Rechunk inputs to make input chunks smaller to avoid an increase in
+        # output chunks
+        new_ops = []
+        for op, input in zip(ops, inputs):
+            changeable_dimensions = {ctr for ctr, i in enumerate(input) if i in outputs}
+            f = max(factor ** (len(changeable_dimensions) / len(outputs)), 1)
+            result = _calculate_new_chunksizes(
+                op.chunks,
+                list(op.chunks),
+                changeable_dimensions,
+                math.prod(map(cached_max, op.chunks)) / f,
+            )
+            new_ops.append(op.rechunk(result))
+        ops = new_ops
+
+    # Introduce the contracted indices into the blockwise product
+    # so that we get numpy arrays, not lists
+    result = blockwise(
+        chunk_einsum,
+        tuple(outputs) + tuple(contract_inds),
+        *(a for ap in zip(ops, inputs) for a in ap),
+        # blockwise parameters
+        adjust_chunks=dict.fromkeys(contract_inds, 1),
+        dtype=dtype,
+        # np.einsum parameters
+        subscripts=subscripts,
+        kernel_dtype=einsum_dtype,
+        ncontract_inds=ncontract_inds,
+        optimize=optimize,
+        **kwargs,
+    )
+
+    # Now reduce over any extra contraction dimensions
+    if ncontract_inds > 0:
+        size = len(outputs)
+        return result.sum(axis=list(range(size, size + ncontract_inds)), split_every=split_every)
+
+    return result