Blosc / Blosc/python-blosc2

Fully lazy imperative function evaluation

Open
#499 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
211
Forks
62
Avg merge
1d 17h
Merged PRs (30d)
6

Description

In order to make function evaluation properly lazy, all functions (e.g. matmul) should be implemented as LazyUDFs.

Secondly, the lazyexpr machinery of compute should loop over chunks of the result. Each function must then decide what slices of the operands are necessary to form the corresponding result chunk (as matmul currently does internally). Then, upon evaluation, although the expression evaluates term-by-term, it does not compute the full result for e..g matmul before proceeding to the next term, but only for the necessary chunk(s) of output. Thus there is a higher chance of cache hits (not the case currently for eager execution of linalg and reductions).
Something in this spirit has been implemented for cumsum/cumprod when axis=None (see this code ):

# Special case for cumulative operations with axis = None
        if reduce_args["axis"] is None and reduce_op in {ReduceOp.CUMULATIVE_PROD, ReduceOp.CUMULATIVE_SUM}:
            # res_out_ is just None, out set to all 0s (sum) or 1s (prod)
            out, res_out_ = convert_none_out(dtype, reduce_op, reduced_shape)
            # reduced_shape is just one-element tuple
            chunklen = out.chunks[0] if hasattr(out, "chunks") else chunks[-1]
            carry = 0
            for cidx in range(0, reduced_shape[0] // chunklen):
                slice_starts = np.unravel_index(cidx * chunklen, shape)
                slice_stops = np.unravel_index((cidx + 1) * chunklen, shape)
                cslice = tuple(
                    slice(start, stop) for start, stop in zip(slice_starts, slice_stops, strict=True)
                )
                _get_chunk_operands(operands, cslice, chunk_operands, shape)
                result, _ = _get_result(expression, chunk_operands, ne_args, where)
                result = np.require(result, requirements="C")
                if reduce_op == ReduceOp.CUMULATIVE_SUM:
                    res = np.cumulative_sum(result, axis=None) + carry
                else:
                    res = np.cumulative_prod(result, axis=None) * carry
                carry = res[-1]
                out[cidx * chunklen + include_initial : (cidx + 1) * chunklen + include_initial] = res

It should also be implemented for fancy-indexing with slice (see https://github.com/Blosc/python-blosc2/issues/441).

Should be possible to handle even something like "matmul(sum(a,axis=1), b)" by passing the desired slice from matmul->sum, which treats the asked-for-slice as a desired output slice and handles accordingly.

This would avoid large in-memory temporaries for example when calculating from eagerly executed linear algebra functions e.g. in "matmul(a, b) + b".

Problems:
1 - reductions could be a problem since for example "sum(a) + a" for the chunks of output would recalculate the scalar "sum(a)" for each chunk of output. This could be avoided by using some kind of result cache for reductions.
2 - naturally, numexpr would always be faster since it compiles the expression into bytecode. Thus we should make sure that, when possible, we still use numexpr preferentially (for most elementwise funcs).

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Read the linked cumsum/cumprod axis=None implementation and trace the lazyexpr compute machinery, LazyUDFs, and matmul chunk handling. Check issue #441 for the fancy-indexing context. Done means functions evaluate requested output chunks lazily, including nested matmul/sum expressions, without full-result temporaries while preserving numexpr preference and avoiding repeated reduction work.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.