From f7d23b4e412fed6cae4d889f99e9573a34213fbc Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Tue, 18 Aug 2026 10:58:31 +0300 Subject: [PATCH 1/4] Numba: generate Join with named parameters and slice writes 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 --- pytensor/link/numba/dispatch/tensor_basic.py | 85 ++++++++++++++++++-- tests/link/numba/test_tensor_basic.py | 8 ++ 2 files changed, 85 insertions(+), 8 deletions(-) diff --git a/pytensor/link/numba/dispatch/tensor_basic.py b/pytensor/link/numba/dispatch/tensor_basic.py index dcaa540d26..68d4c91928 100644 --- a/pytensor/link/numba/dispatch/tensor_basic.py +++ b/pytensor/link/numba/dispatch/tensor_basic.py @@ -9,7 +9,11 @@ register_funcify_and_cache_key, register_funcify_default_op_cache_key, ) -from pytensor.link.numba.dispatch.string_codegen import create_tuple_string +from pytensor.link.numba.dispatch.string_codegen import ( + CODE_TOKEN, + build_source_code, + create_tuple_string, +) from pytensor.tensor.basic import ( Alloc, AllocEmpty, @@ -115,14 +119,79 @@ def arange(start, stop, step): @register_funcify_default_op_cache_key(Join) -def numba_funcify_Join(op, **kwargs): - axis = op.axis - - @numba_basic.numba_njit - def join(*tensors): - return np.concatenate(tensors, axis) +def numba_funcify_Join(op, node, **kwargs): + """Write each input into its slice of a preallocated output. + + ``np.concatenate`` on a tuple of arrays compiles ~3x slower for the same + result. For ``join(1, x0, x1, x2)`` on matrices this emits:: + + def join(*tensors): + total = tensors[0].shape[1] + total += tensors[1].shape[1] + total += tensors[2].shape[1] + if tensors[1].shape[0] != tensors[0].shape[0]: + raise ValueError(_mismatch_msg) + if tensors[2].shape[0] != tensors[0].shape[0]: + raise ValueError(_mismatch_msg) + out = np.empty((tensors[0].shape[0], total), dtype) + off = 0 + l = tensors[0].shape[1] + out[:, off : off + l] = tensors[0] + off += l + ... + return out + """ + ndim = node.outputs[0].type.ndim + ax = op.axis % ndim + pre = ":, " * ax + post = ", :" * (ndim - ax - 1) + names = [f"tensors[{i}]" for i in range(len(node.inputs))] + + shape = [f"tensors[0].shape[{d}]" for d in range(ndim)] + shape[ax] = "total" + code: list[str | CODE_TOKEN] = [ + "def join(*tensors):", + CODE_TOKEN.INDENT, + f"total = tensors[0].shape[{ax}]", + *(f"total += {name}.shape[{ax}]" for name in names[1:]), + ] + for name in names[1:]: + for d in range(ndim): + if d != ax: + code += [ + f"if {name}.shape[{d}] != tensors[0].shape[{d}]:", + CODE_TOKEN.INDENT, + "raise ValueError(_mismatch_msg)", + CODE_TOKEN.DEDENT, + ] + code += [ + f"out = np.empty({create_tuple_string(shape)}, dtype)", + "off = 0", + ] + for name in names: + code += [ + f"l = {name}.shape[{ax}]", + f"out[{pre}off:off + l{post}] = {name}", + "off += l", + ] + code += ["return out", CODE_TOKEN.DEDENT] + + join_fn = compile_numba_function_src( + build_source_code(code), + "join", + globals() + | { + "np": np, + "dtype": np.dtype(node.outputs[0].type.dtype), + "_mismatch_msg": "all the input array dimensions except for the " + "concatenation axis must match exactly", + }, + ) - return join + cache_version = 3 + # The explicit slice writes cannot go out of bounds: the output is + # allocated from the validated input shapes + return numba_basic.numba_njit(join_fn, boundscheck=False), cache_version @register_funcify_default_op_cache_key(Split) diff --git a/tests/link/numba/test_tensor_basic.py b/tests/link/numba/test_tensor_basic.py index 23cf6754fe..c127ba5152 100644 --- a/tests/link/numba/test_tensor_basic.py +++ b/tests/link/numba/test_tensor_basic.py @@ -172,6 +172,14 @@ def test_Join(vals, axis): ) +def test_Join_mismatched_shape_raises(): + xs = [pt.matrix("x0"), pt.matrix("x1")] + fn = function(xs, pt.join(0, *xs), mode=get_mode("NUMBA")) + rng = np.random.default_rng(0) + with pytest.raises(ValueError, match="dimensions except for the concatenation"): + fn(rng.normal(size=(2, 3)), rng.normal(size=(2, 4))) + + @pytest.mark.parametrize( "n_splits, axis, values, sizes", [ From a693ddd8bfa99b81a1bfd89da3a8fff285f80bfe Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Mon, 24 Aug 2026 21:03:52 +0200 Subject: [PATCH 2/4] More performant Join --- pytensor/link/numba/dispatch/tensor_basic.py | 45 +++++++++++++------- tests/link/numba/test_tensor_basic.py | 18 ++++++++ 2 files changed, 48 insertions(+), 15 deletions(-) diff --git a/pytensor/link/numba/dispatch/tensor_basic.py b/pytensor/link/numba/dispatch/tensor_basic.py index 68d4c91928..6e9c66b99b 100644 --- a/pytensor/link/numba/dispatch/tensor_basic.py +++ b/pytensor/link/numba/dispatch/tensor_basic.py @@ -120,10 +120,14 @@ def arange(start, stop, step): @register_funcify_default_op_cache_key(Join) def numba_funcify_Join(op, node, **kwargs): - """Write each input into its slice of a preallocated output. + """Copy each input into its slice of a preallocated output. - ``np.concatenate`` on a tuple of arrays compiles ~3x slower for the same - result. For ``join(1, x0, x1, x2)`` on matrices this emits:: + ``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 + rebuilds the source as an ``A``-layout view and so copies it with gathers; the + dimensions ahead of the join axis are peeled with integer indexing so that the + destination slice stays contiguous. For ``join(1, x0, x1, x2)`` on matrices this + emits:: def join(*tensors): total = tensors[0].shape[1] @@ -136,15 +140,17 @@ def join(*tensors): out = np.empty((tensors[0].shape[0], total), dtype) off = 0 l = tensors[0].shape[1] - out[:, off : off + l] = tensors[0] + for i0 in range(tensors[0].shape[0]): + dst = out[i0][off : off + l] + src = tensors[0][i0] + for j0 in range(src.shape[0]): + dst[j0] = src[j0] off += l ... return out """ ndim = node.outputs[0].type.ndim - ax = op.axis % ndim - pre = ":, " * ax - post = ", :" * (ndim - ax - 1) + ax = op.axis names = [f"tensors[{i}]" for i in range(len(node.inputs))] shape = [f"tensors[0].shape[{d}]" for d in range(ndim)] @@ -168,12 +174,21 @@ def join(*tensors): f"out = np.empty({create_tuple_string(shape)}, dtype)", "off = 0", ] + # Dimensions before the join axis are peeled with integer indexing, so that the + # destination slice on the join axis stays contiguous whenever the input is. + leading = "".join(f"[i{d}]" for d in range(ax)) + trailing_idxs = ", ".join(f"j{k}" for k in range(ndim - ax)) for name in names: - code += [ - f"l = {name}.shape[{ax}]", - f"out[{pre}off:off + l{post}] = {name}", - "off += l", - ] + code += [f"l = {name}.shape[{ax}]"] + for d in range(ax): + code += [f"for i{d} in range({name}.shape[{d}]):", CODE_TOKEN.INDENT] + code += [f"dst = out{leading}[off:off + l]", f"src = {name}{leading}"] + for k in range(ndim - ax): + code += [f"for j{k} in range(src.shape[{k}]):", CODE_TOKEN.INDENT] + code += [f"dst[{trailing_idxs}] = src[{trailing_idxs}]"] + code += [CODE_TOKEN.DEDENT] * (ndim - ax) + code += [CODE_TOKEN.DEDENT] * ax + code += ["off += l"] code += ["return out", CODE_TOKEN.DEDENT] join_fn = compile_numba_function_src( @@ -188,9 +203,9 @@ def join(*tensors): }, ) - cache_version = 3 - # The explicit slice writes cannot go out of bounds: the output is - # allocated from the validated input shapes + cache_version = 4 + # The loop nests cannot go out of bounds: the output is allocated from the + # validated input shapes return numba_basic.numba_njit(join_fn, boundscheck=False), cache_version diff --git a/tests/link/numba/test_tensor_basic.py b/tests/link/numba/test_tensor_basic.py index c127ba5152..152f711c79 100644 --- a/tests/link/numba/test_tensor_basic.py +++ b/tests/link/numba/test_tensor_basic.py @@ -159,6 +159,24 @@ def test_ARange(): ), 1, ), + # More than two inputs, ragged along the join axis, with dimensions both + # ahead of and behind it + ( + ( + (pt.tensor3(), rng.normal(size=(2, 1, 3)).astype(config.floatX)), + (pt.tensor3(), rng.normal(size=(2, 2, 3)).astype(config.floatX)), + (pt.tensor3(), rng.normal(size=(2, 3, 3)).astype(config.floatX)), + ), + 1, + ), + ( + ( + (pt.tensor3(), rng.normal(size=(2, 3, 1)).astype(config.floatX)), + (pt.tensor3(), rng.normal(size=(2, 3, 2)).astype(config.floatX)), + (pt.tensor3(), rng.normal(size=(2, 3, 3)).astype(config.floatX)), + ), + 2, + ), ], ) def test_Join(vals, axis): From 96e6df66ed9a090021b954ba7fb913b5cf4c4769 Mon Sep 17 00:00:00 2001 From: Margus Niitsoo Date: Tue, 25 Aug 2026 11:35:24 +0300 Subject: [PATCH 3/4] Single-step destination indexing; full codegen example in the docstring 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) --- pytensor/link/numba/dispatch/tensor_basic.py | 34 ++++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/pytensor/link/numba/dispatch/tensor_basic.py b/pytensor/link/numba/dispatch/tensor_basic.py index 6e9c66b99b..5bcec7b6a9 100644 --- a/pytensor/link/numba/dispatch/tensor_basic.py +++ b/pytensor/link/numba/dispatch/tensor_basic.py @@ -122,11 +122,13 @@ def arange(start, stop, step): def numba_funcify_Join(op, node, **kwargs): """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 - rebuilds the source as an ``A``-layout view and so copies it with gathers; the - dimensions ahead of the join axis are peeled with integer indexing so that the - destination slice stays contiguous. For ``join(1, x0, x1, x2)`` on matrices this + ``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 rebuilds the source as an ``A``-layout view and so copies it + with gathers; the dimensions ahead of the join axis are peeled with integer + indexing so that the destination slice stays contiguous. + + For ``join(1, x0, x1, x2)`` on matrices (axis 1, the last one here) this emits:: def join(*tensors): @@ -141,12 +143,25 @@ def join(*tensors): off = 0 l = tensors[0].shape[1] for i0 in range(tensors[0].shape[0]): - dst = out[i0][off : off + l] + dst = out[i0, off : off + l] src = tensors[0][i0] for j0 in range(src.shape[0]): dst[j0] = src[j0] off += l - ... + l = tensors[1].shape[1] + for i0 in range(tensors[1].shape[0]): + dst = out[i0, off : off + l] + src = tensors[1][i0] + for j0 in range(src.shape[0]): + dst[j0] = src[j0] + off += l + l = tensors[2].shape[1] + for i0 in range(tensors[2].shape[0]): + dst = out[i0, off : off + l] + src = tensors[2][i0] + for j0 in range(src.shape[0]): + dst[j0] = src[j0] + off += l return out """ ndim = node.outputs[0].type.ndim @@ -174,15 +189,14 @@ def join(*tensors): f"out = np.empty({create_tuple_string(shape)}, dtype)", "off = 0", ] - # Dimensions before the join axis are peeled with integer indexing, so that the - # destination slice on the join axis stays contiguous whenever the input is. leading = "".join(f"[i{d}]" for d in range(ax)) + leading_idx = "".join(f"i{d}, " for d in range(ax)) trailing_idxs = ", ".join(f"j{k}" for k in range(ndim - ax)) for name in names: code += [f"l = {name}.shape[{ax}]"] for d in range(ax): code += [f"for i{d} in range({name}.shape[{d}]):", CODE_TOKEN.INDENT] - code += [f"dst = out{leading}[off:off + l]", f"src = {name}{leading}"] + code += [f"dst = out[{leading_idx}off:off + l]", f"src = {name}{leading}"] for k in range(ndim - ax): code += [f"for j{k} in range(src.shape[{k}]):", CODE_TOKEN.INDENT] code += [f"dst[{trailing_idxs}] = src[{trailing_idxs}]"] From 42a695ac8a9182df2a39e8b54cf44f81cd5a9019 Mon Sep 17 00:00:00 2001 From: ricardoV94 Date: Mon, 31 Aug 2026 16:52:33 +0200 Subject: [PATCH 4/4] Numba: emit Join as a scalar loop nest with an unsigned offset One flat scalar nest per input, with the running offset added to the join axis subscript, replaces the per-input slice write and its peeled destination views. The offset is unsigned so LLVM can prove the destination index non-negative; with a signed offset it vectorises only the first input's copy loop. Broadcastable dimensions are indexed with a constant instead of looped over, so that the innermost loop -- the one LLVM vectorises -- is one that varies. --- pytensor/link/numba/dispatch/tensor_basic.py | 78 +++++++++----------- tests/link/numba/test_tensor_basic.py | 14 ++++ 2 files changed, 47 insertions(+), 45 deletions(-) diff --git a/pytensor/link/numba/dispatch/tensor_basic.py b/pytensor/link/numba/dispatch/tensor_basic.py index 5bcec7b6a9..2993c42862 100644 --- a/pytensor/link/numba/dispatch/tensor_basic.py +++ b/pytensor/link/numba/dispatch/tensor_basic.py @@ -120,16 +120,13 @@ def arange(start, stop, step): @register_funcify_default_op_cache_key(Join) def numba_funcify_Join(op, node, **kwargs): - """Copy each input into its slice of a preallocated output. + """Copy each input into a preallocated output with a scalar loop nest. - ``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 rebuilds the source as an ``A``-layout view and so copies it - with gathers; the dimensions ahead of the join axis are peeled with integer - indexing so that the destination slice stays contiguous. + ``np.concatenate`` on a tuple of arrays compiles slower for the same result. + To let LLVM vectorize the inner copy loop, the join axis offset is unsigned + (provably non-negative) and broadcastable dimensions are indexed with zero. - For ``join(1, x0, x1, x2)`` on matrices (axis 1, the last one here) this - emits:: + For ``join(1, x0, x1, x2)`` on matrices this emits:: def join(*tensors): total = tensors[0].shape[1] @@ -140,28 +137,19 @@ def join(*tensors): if tensors[2].shape[0] != tensors[0].shape[0]: raise ValueError(_mismatch_msg) out = np.empty((tensors[0].shape[0], total), dtype) - off = 0 - l = tensors[0].shape[1] - for i0 in range(tensors[0].shape[0]): - dst = out[i0, off : off + l] - src = tensors[0][i0] - for j0 in range(src.shape[0]): - dst[j0] = src[j0] - off += l - l = tensors[1].shape[1] - for i0 in range(tensors[1].shape[0]): - dst = out[i0, off : off + l] - src = tensors[1][i0] - for j0 in range(src.shape[0]): - dst[j0] = src[j0] - off += l - l = tensors[2].shape[1] - for i0 in range(tensors[2].shape[0]): - dst = out[i0, off : off + l] - src = tensors[2][i0] - for j0 in range(src.shape[0]): - dst[j0] = src[j0] - off += l + off = np.uint64(0) + for j0 in range(tensors[0].shape[0]): + for j1 in range(tensors[0].shape[1]): + out[j0, np.uint64(j1) + off] = tensors[0][j0, j1] + off += np.uint64(tensors[0].shape[1]) + for j0 in range(tensors[1].shape[0]): + for j1 in range(tensors[1].shape[1]): + out[j0, np.uint64(j1) + off] = tensors[1][j0, j1] + off += np.uint64(tensors[1].shape[1]) + for j0 in range(tensors[2].shape[0]): + for j1 in range(tensors[2].shape[1]): + out[j0, np.uint64(j1) + off] = tensors[2][j0, j1] + off += np.uint64(tensors[2].shape[1]) return out """ ndim = node.outputs[0].type.ndim @@ -187,22 +175,22 @@ def join(*tensors): ] code += [ f"out = np.empty({create_tuple_string(shape)}, dtype)", - "off = 0", + "off = np.uint64(0)", ] - leading = "".join(f"[i{d}]" for d in range(ax)) - leading_idx = "".join(f"i{d}, " for d in range(ax)) - trailing_idxs = ", ".join(f"j{k}" for k in range(ndim - ax)) + # The join axis is never broadcastable here: its static length is the sum. + static_shape = node.outputs[0].type.shape + looped = [d for d in range(ndim) if d == ax or static_shape[d] != 1] + src_idx = ", ".join(f"j{d}" if d in looped else "0" for d in range(ndim)) + dst_idx = ", ".join( + f"np.uint64(j{d}) + off" if d == ax else (f"j{d}" if d in looped else "0") + for d in range(ndim) + ) for name in names: - code += [f"l = {name}.shape[{ax}]"] - for d in range(ax): - code += [f"for i{d} in range({name}.shape[{d}]):", CODE_TOKEN.INDENT] - code += [f"dst = out[{leading_idx}off:off + l]", f"src = {name}{leading}"] - for k in range(ndim - ax): - code += [f"for j{k} in range(src.shape[{k}]):", CODE_TOKEN.INDENT] - code += [f"dst[{trailing_idxs}] = src[{trailing_idxs}]"] - code += [CODE_TOKEN.DEDENT] * (ndim - ax) - code += [CODE_TOKEN.DEDENT] * ax - code += ["off += l"] + for d in looped: + code += [f"for j{d} in range({name}.shape[{d}]):", CODE_TOKEN.INDENT] + code += [f"out[{dst_idx}] = {name}[{src_idx}]"] + code += [CODE_TOKEN.DEDENT] * len(looped) + code += [f"off += np.uint64({name}.shape[{ax}])"] code += ["return out", CODE_TOKEN.DEDENT] join_fn = compile_numba_function_src( @@ -217,7 +205,7 @@ def join(*tensors): }, ) - cache_version = 4 + cache_version = 5 # The loop nests cannot go out of bounds: the output is allocated from the # validated input shapes return numba_basic.numba_njit(join_fn, boundscheck=False), cache_version diff --git a/tests/link/numba/test_tensor_basic.py b/tests/link/numba/test_tensor_basic.py index 152f711c79..a41238f2ee 100644 --- a/tests/link/numba/test_tensor_basic.py +++ b/tests/link/numba/test_tensor_basic.py @@ -177,6 +177,20 @@ def test_ARange(): ), 2, ), + # A broadcastable dimension is indexed with a constant rather than looped over + ( + ( + ( + pt.tensor("x0", shape=(None, None, 1)), + rng.normal(size=(2, 1, 1)).astype(config.floatX), + ), + ( + pt.tensor("x1", shape=(None, None, 1)), + rng.normal(size=(2, 3, 1)).astype(config.floatX), + ), + ), + 1, + ), ], ) def test_Join(vals, axis):