pymc-devs / pymc-devs/pytensor

Inplace Blockwise outputs on inputs in Numba backend

Open
#1,030 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

numba Op implementation performance vectorization
Dominant language
Python
Stars
644
Forks
208
Avg merge
2d 14h
Merged PRs (30d)
16

Description

Description

When we have a Blockwise with the same output shape and type as one of its inputs (such as Cholesky whose signature is "(m,m)->(m,m)") we can overwrite the output in the input buffer.

This is independent of allowing the core_op to destroy the input as done in #1028

This is only relevant for the numba backend after #1015 since the other backends don't give us inplace control and numpy vectorize (used in the Python backend) neither.

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 Numba backend's Blockwise handling first, using the Cholesky signature and the distinctions from #1015 and #1028 as context. Determine how same-shaped, same-typed outputs are allocated and verify that the completed behavior overwrites the input buffer only for the Numba backend, without changing the other backends.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.