Reuse output of Join in C backend

pytensor/link/jax/dispatch/tensor_basic.py

@@ -87,14 +87,7 @@ def jax_funcify_Join(op, **kwargs):
     def join(axis, *tensors):
         # tensors could also be tuples, and in this case they don't have a ndim
         tensors = [jnp.asarray(tensor) for tensor in tensors]
-        view = op.view
-        if (view != -1) and all(
-            tensor.shape[axis] == 0 for tensor in tensors[0:view] + tensors[view + 1 :]
-        ):
-            return tensors[view]
-
-        else:
-            return jnp.concatenate(tensors, axis=axis)
+        return jnp.concatenate(tensors, axis=axis)
 
     return join
 

pytensor/link/numba/dispatch/tensor_basic.py

@@ -117,17 +117,9 @@ def arange(start, stop, step):
 
 @numba_funcify.register(Join)
 def numba_funcify_Join(op, **kwargs):
-    view = op.view
-
-    if view != -1:
-        # TODO: Where (and why) is this `Join.view` even being used?  From a
-        # quick search, the answer appears to be "nowhere", so we should
-        # probably just remove it.
-        raise NotImplementedError("The `view` parameter to `Join` is not supported")
-
     @numba_basic.numba_njit
     def join(axis, *tensors):
-        return np.concatenate(tensors, numba_basic.to_scalar(axis))
+        return np.concatenate(tensors, axis.item())
 
     return join
 

pytensor/scan/checkpoints.py

@@ -1,6 +1,5 @@
 import pytensor.tensor.basic as ptb
 from pytensor.scan.basic import scan
-from pytensor.tensor.basic import Join
 from pytensor.tensor.math import ceil, eq, neq
 from pytensor.tensor.subtensor import set_subtensor
 
@@ -127,14 +126,12 @@ def scan_checkpoints(
 
     # Pad the sequences if needed
     if padding:
-        # Since padding could be an empty tensor, Join returns a view of s.
-        join = Join(view=0)
         for i, s in enumerate(sequences):
             overshoots_by = s.shape[0] % save_every_N
             overshoots = neq(overshoots_by, 0)
             n = (save_every_N - overshoots_by) * overshoots
             z = ptb.zeros((n, *s.shape[1:]), dtype=s.dtype)
-            sequences[i] = join(0, s, z)
+            sequences[i] = ptb.join(0, s, z)
 
     # Establish the input variables of the outer scan
     o_sequences = [