diff --git a/pytensor/link/numba/dispatch/tensor_basic.py b/pytensor/link/numba/dispatch/tensor_basic.py index dcaa540d26..2993c42862 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,96 @@ 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): + """Copy each input into a preallocated output with a scalar loop nest. + + ``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 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 = 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 + 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)] + 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 = np.uint64(0)", + ] + # 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: + 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( + 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 = 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 @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..a41238f2ee 100644 --- a/tests/link/numba/test_tensor_basic.py +++ b/tests/link/numba/test_tensor_basic.py @@ -159,6 +159,38 @@ 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, + ), + # 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): @@ -172,6 +204,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", [