Numba: copy Join inputs into a preallocated output instead of np.concatenate by velochy · Pull Request #2362 · pymc-devs/pytensor

velochy · 2026-08-18T11:46:15Z

Join's numba implementation was np.concatenate on a *tensors tuple; numba types and lowers an implementation that receives its inputs as one wide tuple quadratically in the input count, and >30-input joins are routine (gradient assembly concatenates one piece per parameter).

The generated implementation takes one named parameter per input and writes each into its slice of a preallocated output — linear codegen, same single copy per input. Split off from #2354 per review there (a dispatch-level Join fix instead of a graph rewrite).

Results

Interleaved with main in one session on an otherwise-idle 8-core/16-thread box (Ryzen 7 PRO 4750U, BLAS/OMP threads pinned to 1), cold caches (absolute times here are load-sensitive, so only same-session comparisons are quoted), n=80 additive repro whose gradient assembly is a 60-tensor join:

	compile	peak RSS
main	364.9 s	2696 MB
this PR alone	151.5 s	2660 MB
with #2361	105.9 s	933 MB

The two are complementary: this one removes the wide-tuple callee (np.concatenate typed with a 60-tuple) and carries the compile-time win, while #2361 removes the caller-side bytecode chain and carries the memory win.

ricardoV94 · 2026-08-18T13:05:53Z

+    post = ", :" * (ndim - ax - 1)
+    names = [f"x{i}" for i in range(len(node.inputs))]
+
+    total_src = "\n".join(f"    total += {name}.shape[{ax}]" for name in names[1:])


can you use our string codegen helpers? makes it a tiny bit more readable

Done — rewritten with CODE_TOKEN/build_source_code.

ricardoV94 · 2026-08-18T13:07:19Z


-    return join
+    cache_version = 3
+    return numba_basic.numba_njit(join_fn), cache_version


checkbounds=False iff the function already checks shapes are correct

Added boundscheck=False — the slice writes are within the output allocated from the validated shapes by construction.

velochy · 2026-08-23T20:17:44Z

You were right about the argument names, so I removed them: the signature is back to *tensors and the body indexes it. I isolated the two changes on top of #2361 — np.concatenate with a *tensors bundle 325.6 / 323.7 s, slice-writes with the same bundle 107.0 / 107.1 s, slice-writes with named parameters 107.1 / 106.3 s — so the entire 3x is the preallocate-and-slice-write body and the names are worth 0%, exactly as in #2354. Consistent rather than contradictory: np.concatenate is a callee that genuinely consumes the tuple as a tuple, while the slice-write body never touches it as a whole. The shape validation and boundscheck=False stay.

ricardoV94 · 2026-08-23T23:17:59Z

    )


+@pytest.mark.parametrize("n_terms, axis", [(35, 0), (35, 1), (35, -1)])


are the new tests relevant now? Sounds like no?

Right on both counts. The MakeVector one was a leftover from when this PR was about wide calls — MakeVector's dispatch isn't touched here at all, so it's gone. The 35-input Join case was testing arity, which stopped being a distinct code path once the win turned out to be the slice-write body rather than the signature; the existing test_Join already covers the shape of it.

What is genuinely new is the running offset across inputs and the per-input shape checks, so I replaced both with a small test_Join_multiple_ragged (four inputs, ragged along the join axis, axis 0/1/-1 — three or more inputs is what makes the offset accumulate more than once) plus a direct check that mismatched non-join dims raise ValueError. Faster than what it replaces and it actually exercises the new code.

ricardoV94 · 2026-08-24T09:42:04Z



+@pytest.mark.parametrize("axis", [0, 1, -1])
+def test_Join_multiple_ragged(axis):


I still don't think this is new/needed, it already tested above/elsewhere in the standard suite. You can leave the one with error which I'm less sure was tested comprehensively

Added — the docstring now carries the emitted source for join(1, x0, x1, x2) on matrices, copied from what the dispatcher actually generates (shape checks, the preallocated out, and the running offset), in the same style as store_core_outputs.

ricardoV94 · 2026-08-24T09:42:47Z

-    @numba_basic.numba_njit
-    def join(*tensors):
-        return np.concatenate(tensors, axis)
+def numba_funcify_Join(op, node, **kwargs):


can you include an example of the final codegen for a real join in the docstrings. I did that in some other dispatchers, makes it easier to grasp what sort of code is being emitted.

Dropped it — you're right, TestJoinAndSplit already covers multi-input and negative-axis joins and runs in NUMBA mode in CI (7 cases with three or more inputs, 8 with a negative axis), so it was duplicate coverage. Kept only the mismatched-shape error test.

np.concatenate typed with an n-tuple lowers as O(n^2) LLVM IR; one named parameter and one slice-write per input into a preallocated output is linear, with the same single copy per input. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

ricardoV94 · 2026-08-25T06:50:49Z

@velochy it looked like your implementation had worse runtime performance. I tried to push a version with explicit loop that seems similar to np.concatenate in perf while still compiling faster.

Can you check it? The PR title needs a change as it still mentions the input name thing.

I haven't had the time to push against some of the specofics my commit came up with

ricardoV94 · 2026-08-25T06:57:24Z

+    """Copy each input into its slice of a preallocated output.
+
+    ``np.concatenate`` on a tuple of arrays compiles ~3x slower for the same result.
+    A whole-array slice write instead goes through Numba's fancy indexing, which


I don't like the comment my bot added. Also I'd rather have the whole snippet and no ....

And an explicit mention of what's the axis so the reader doesn't have to guess it

Removed the comment, and the docstring now carries the whole emitted function with no elision and says which axis it is showing (axis 1, the last one for matrices).

ricardoV94 · 2026-08-25T06:58:27Z

+            off = 0
+            l = tensors[0].shape[1]
+            for i0 in range(tensors[0].shape[0]):
+                dst = out[i0][off : off + l]


I'm surprised if we really need this two step indexing for perf...

You were right to be surprised — it isn't needed. Timing the four codegen shapes directly (mean of 3 runs, each best-of-60 calls, contiguous inputs):

geometry np.concatenate whole-array slice write loops, out[i0][off:off+l] loops, out[i0, off:off+l]

n=60, 1-D, axis 0 11.6 us 11.2 11.5 11.2

n=8, 2-D, axis 1 13.4 us 19.4 (+45%) 15.9 13.7

n=8, 2-D, axis 0 9.1 us 15.2 (+67%) 9.2 9.3

n=6, 3-D, axis 2 (62 MB) 45.2 ms 50.1 (+11%) 45.2 45.3

Single-step is at worst equal, so I switched to it — the raw samples also make it look the steadier of the two (13.3/14.0/13.7 and 9.4/9.0/9.6, against 13.8/13.5/20.4 and 8.1/11.5/8.1 for the two-step form; those excursions are what push its means up in the table).

out[i0, off:off + l] measures the same as out[i0][off:off + l] on every geometry tried (2-D axis 1: 13.7 vs 15.9 us, 2-D axis 0: 9.3 vs 9.2 us, 3-D: 45.3 vs 45.2 ms), so keep the simpler form. The docstring now carries the whole emitted function and names the axis it is showing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

velochy · 2026-08-25T08:35:57Z

Retitled, and your runtime observation was correct — the whole-array slice write really was slower, and only on multi-dimensional joins, which is why my model-level benchmarks never saw it (the gradient-assembly join that motivated the PR is 1-D, where all variants tie at ~11 us). Measured per-call, mean of 3 runs each best-of-60: 2-D axis 1 concat 13.4 us vs slice 19.4 (+45%), 2-D axis 0 9.1 vs 15.2 (+67%), 3-D 45.2 ms vs 50.1 (+11%). Your loop version restores parity everywhere, so it stays; I only simplified the destination indexing to a single step since it measures the same (see the inline reply).

I have a full riigikogu26 fit running for all three variants to confirm compile time and memory stay where they were — will post those numbers when it finishes.

velochy mentioned this pull request Aug 18, 2026

Numba: rework the vectorize intrinsic interface to fix wide-node compile blowup #2354

Open

velochy force-pushed the join-dispatch branch from d95914f to f1e413d Compare August 18, 2026 12:51

ricardoV94 reviewed Aug 18, 2026

View reviewed changes

velochy force-pushed the join-dispatch branch from f1e413d to ed90e10 Compare August 18, 2026 13:15

velochy mentioned this pull request Aug 20, 2026

Numba: work around quadratic lowering of >30-argument calls #2361

Open

velochy force-pushed the join-dispatch branch from ed90e10 to 21f9157 Compare August 23, 2026 20:17

ricardoV94 reviewed Aug 23, 2026

View reviewed changes

velochy force-pushed the join-dispatch branch from 21f9157 to f67494e Compare August 24, 2026 07:36

ricardoV94 reviewed Aug 24, 2026

View reviewed changes

velochy force-pushed the join-dispatch branch from f67494e to f7d23b4 Compare August 24, 2026 12:15

More performant Join

a693ddd

ricardoV94 reviewed Aug 25, 2026

View reviewed changes

velochy changed the title ~~Numba: generate Join with named parameters and slice writes~~ Numba: copy Join inputs into a preallocated output instead of np.concatenate Aug 25, 2026

		)


		@pytest.mark.parametrize("n_terms, axis", [(35, 0), (35, 1), (35, -1)])



		@pytest.mark.parametrize("axis", [0, 1, -1])
		def test_Join_multiple_ragged(axis):

geometry	`np.concatenate`	whole-array slice write	loops, `out[i0][off:off+l]`	loops, `out[i0, off:off+l]`
n=60, 1-D, axis 0	11.6 us	11.2	11.5	11.2
n=8, 2-D, axis 1	13.4 us	19.4 (+45%)	15.9	13.7
n=8, 2-D, axis 0	9.1 us	15.2 (+67%)	9.2	9.3
n=6, 3-D, axis 2 (62 MB)	45.2 ms	50.1 (+11%)	45.2	45.3

Uh oh!

Conversation

velochy commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Results

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

velochy commented Aug 23, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

velochy commented Aug 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

velochy commented Aug 18, 2026 •

edited

Loading

ricardoV94 commented Aug 25, 2026 •

edited

Loading