Repository: coreylowman/dfdx (language: rust)

Issue:
Concat along any dimension
Currently concat only works along the 0th dimension. However it is very common to concat along any dimension. This happens in kv-cache implementations (concat along 2nd dimension), and many others.


Discussion / hints:

Fix patch:
diff --git a/src/tensor_ops/concat/mod.rs b/src/tensor_ops/concat/mod.rs
index 41668d92c..51ef0df64 100644
--- a/src/tensor_ops/concat/mod.rs
+++ b/src/tensor_ops/concat/mod.rs
@@ -26,18 +26,25 @@ mod cuda_kernel;
 /// let c: Tensor<(usize, Const<3>), f32, _> = a.concat(b);
 /// assert_eq!(c.shape().0, 6);
 /// ```
+#[deprecated = "Use TryConcatAlong instead"]
 pub trait TryConcat<Rhs>: HasErr {
     type Output;
 
     /// Concatenate two tensors along the first dimension.
+    #[deprecated = "Use TryConcatAlong::concat_along instead"]
+    #[allow(deprecated)]
     fn concat(self, rhs: Rhs) -> Self::Output {
+        #[allow(deprecated)]
         self.try_concat(rhs).unwrap()
     }
 
     /// Fallible version of [TryConcat::concat].
+    #[deprecated = "Use TryConcatAlong::try_concat_along instead"]
+    #[allow(deprecated)]
     fn try_concat(self, rhs: Rhs) -> Result<Self::Output, Self::Err>;
 }
 
+#[allow(deprecated)]
 impl<A: Shape, B: Shape, T, R, E: Dtype, D: ConcatKernel<E>> TryConcat<Tensor<B, E, D, R>>
     for Tensor<A, E, D, T>
 where
@@ -46,6 +53,7 @@ where
     R: Tape<E, D>,
 {
     type Output = Tensor<A::Catted, E, D, T>;
+    #[allow(deprecated)]
     fn try_concat(self, rhs: Tensor<B, E, D, R>) -> Result<Self::Output, Self::Err> {
         assert_eq!(
             self.strides,
@@ -137,6 +145,7 @@ where
 }
 
 #[cfg(test)]
+#[allow(deprecated)]
 mod tests {
     use super::*;
     use crate::{tensor_ops::*, tests::*};
diff --git a/src/tensor_ops/concat_along/cpu_kernel.rs b/src/tensor_ops/concat_along/cpu_kernel.rs
new file mode 100644
index 000000000..2d480b954
--- /dev/null
+++ b/src/tensor_ops/concat_along/cpu_kernel.rs
@@ -0,0 +1,72 @@
+use crate::{
+    shapes::*,
+    tensor::{cpu::NdIndex, *},
+};
+
+impl<E: Dtype> super::ConcatAlongKernel<E> for Cpu {
+    fn forward<A: Shape, B: Shape, C: Shape>(
+        &self,
+        ax: usize,
+        a: &Tensor<A, E, Self>,
+        b: &Tensor<B, E, Self>,
+        c: &mut Tensor<C, E, Self>,
+    ) -> Result<(), Self::Err> {
+        let mut a_idx = NdIndex::new(a.shape, a.strides);
+        let mut b_idx = NdIndex::new(b.shape, b.strides);
+
+        let mut a_n = 1;
+        let mut b_n = 1;
+        for i in ax..A::NUM_DIMS {
+            a_n *= a_idx.shape[i];
+            b_n *= b_idx.shape[i];
+        }
+
+        let mut i = 0;
+        let n = c.data.len();
+        let buf = std::sync::Arc::get_mut(&mut c.data).unwrap();
+        while i < n {
+            for _ in 0..a_n {
+                buf[i] = a.data[a_idx.next().unwrap()];
+                i += 1;
+            }
+            for _ in 0..b_n {
+                buf[i] = b.data[b_idx.next().unwrap()];
+                i += 1;
+            }
+        }
+        Ok(())
+    }
+    fn backward<A: Shape, B: Shape>(
+        &self,
+        ax: usize,
+        a: &GhostTensor<A, E, Self>,
+        grad_a: &mut Self::Vec<E>,
+        b: &GhostTensor<B, E, Self>,
+        grad_b: &mut Self::Vec<E>,
+        grad_out: &Self::Vec<E>,
+    ) -> Result<(), Self::Err> {
+        let mut a_idx = NdIndex::new(a.shape, a.strides);
+        let mut b_idx = NdIndex::new(b.shape, b.strides);
+
+        let mut a_n = 1;
+        let mut b_n = 1;
+        for i in ax..A::NUM_DIMS {
+            a_n *= a_idx.shape[i];
+            b_n *= b_idx.shape[i];
+        }
+
+        let mut i = 0;
+        let n = grad_out.len();
+        while i < n {
+            for _ in 0..a_n {
+                grad_a[a_idx.next().unwrap()] += grad_out[i];
+                i += 1;
+            }
+            for _ in 0..b_n {
+                grad_b[b_idx.next().unwrap()] += grad_out[i];
+                i += 1;
+            }
+        }
+        Ok(())
+    }
+}
diff --git a/src/tensor_ops/concat_along/cuda_kernel.rs b/src/tensor_ops/concat_along/cuda_kernel.rs
new file mode 100644
index 000000000..b80a6871f
--- /dev/null
+++ b/src/tensor_ops/concat_along/cuda_kernel.rs
@@ -0,0 +1,190 @@
+use crate::{
+    shapes::*,
+    tensor::{launch_cfg, Cuda, GhostTensor, Tensor},
+};
+use cudarc::{
+    driver::{DeviceSlice, LaunchAsync},
+    nvrtc::{compile_ptx_with_opts, CompileOptions},
+    types::CudaTypeName,
+};
+
+impl<E: Dtype + CudaTypeName> super::ConcatAlongKernel<E> for Cuda {
+    fn forward<A: Shape, B: Shape, C: Shape>(
+        &self,
+        ax: usize,
+        a: &Tensor<A, E, Self>,
+        b: &Tensor<B, E, Self>,
+        c: &mut Tensor<C, E, Self>,
+    ) -> Result<(), Self::Err> {
+        let module_name = std::format!("concat_{}", E::NAME);
+        if !self.dev.has_func(&module_name, "fwd") {
+            let src = KERNEL.replace("$Ty", E::NAME);
+            let opts = CompileOptions {
+                arch: Some(env!("CUDA_COMPUTE_CAP")),
+                include_paths: vec![
+                    env!("CUDA_INCLUDE_DIR").to_string(),
+                    env!("OUT_DIR").to_string(),
+                ],
+                ..Default::default()
+            };
+            let ptx = compile_ptx_with_opts(src, opts).unwrap();
+            self.dev.load_ptx(ptx, &module_name, &["fwd", "bwd"])?;
+        }
+
+        let fwd = self.dev.get_func(&module_name, "fwd").unwrap();
+        let cfg = launch_cfg::<128>(c.data.len() as u32);
+
+        let mut info = Vec::with_capacity(3 + 4 * A::NUM_DIMS);
+        info.push(c.data.len());
+        info.push(A::NUM_DIMS);
+        info.push(ax);
+        info.extend(a.shape.concrete());
+        info.extend(a.strides);
+        info.extend(b.shape.concrete());
+        info.extend(b.strides);
+        let info = self.dev.htod_copy(info)?;
+
+        unsafe {
+            fwd.launch(
+                cfg,
+                (
+                    &info,
+                    a.data.as_ref(),
+                    b.data.as_ref(),
+                    std::sync::Arc::get_mut(&mut c.data).unwrap(),
+                ),
+            )
+        }?;
+
+        Ok(())
+    }
+
+    fn backward<A: Shape, B: Shape>(
+        &self,
+        ax: usize,
+        a: &GhostTensor<A, E, Self>,
+        grad_a: &mut Self::Vec<E>,
+        b: &GhostTensor<B, E, Self>,
+        grad_b: &mut Self::Vec<E>,
+        grad_out: &Self::Vec<E>,
+    ) -> Result<(), Self::Err> {
+        let module_name = std::format!("concat_{}", E::NAME);
+        let bwd = self.dev.get_func(&module_name, "bwd").unwrap();
+        let cfg = launch_cfg::<128>(grad_out.data.len() as u32);
+
+        let mut info = Vec::with_capacity(3 + 4 * A::NUM_DIMS);
+        info.push(grad_out.data.len());
+        info.push(A::NUM_DIMS);
+        info.push(ax);
+        info.extend(a.shape.concrete());
+        info.extend(a.strides);
+        info.extend(b.shape.concrete());
+        info.extend(b.strides);
+        let info = self.dev.htod_copy(info)?;
+
+        unsafe { bwd.launch(cfg, (&info, grad_a, grad_b, grad_out)) }?;
+        Ok(())
+    }
+}
+
+const KERNEL: &str = "
+#include \"cuda_utils.cuh\"
+
+extern \"C\" __global__ void fwd(
+    const size_t *info, // numel, num_dims, axis, a_dims, a_strides, b_dims, b_strides
+    const $Ty *lhs,
+    const $Ty *rhs,
+    $Ty *out
+) {
+    const size_t numel = info[0];
+    const size_t num_dims = info[1];
+    const size_t axis = info[2];
+    const size_t *lhs_dims = info + 3;
+    const size_t *lhs_strides = info + 3 + num_dims;
+    const size_t *rhs_dims = info + 3 + 2 * num_dims;
+    const size_t *rhs_strides = info + 3 + 3 * num_dims;
+
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= numel) { return; }
+
+    // out_dims will be (..., lhs_dims[ax] + rhs_dims[ax], ...)
+
+    // striding lhs & rhs up to the concat'd axis
+    size_t i_tmp = i;
+    size_t lhs_i = 0;
+    size_t rhs_i = 0;
+    for (int d = num_dims - 1; d > axis; d--) {
+        size_t dim_i = i_tmp % lhs_dims[d];
+        lhs_i += dim_i * lhs_strides[d];
+        rhs_i += dim_i * rhs_strides[d];
+        i_tmp /= lhs_dims[d];
+    }
+
+    // figure out if we are using lhs or rhs for this `i`
+    size_t i_along_axis = i_tmp % (lhs_dims[axis] + rhs_dims[axis]);
+    i_tmp /= (lhs_dims[axis] + rhs_dims[axis]);
+
+    // striding lhs & rhs along the rest of the axes
+    for (int d = axis - 1; d >= 0;d--) {
+        size_t dim_i = i_tmp % lhs_dims[d];
+        lhs_i += dim_i * lhs_strides[d];
+        rhs_i += dim_i * rhs_strides[d];
+        i_tmp /= lhs_dims[d];
+    }
+
+    if (i_along_axis < lhs_dims[axis]) {
+        out[i] = lhs[lhs_i + i_along_axis * lhs_strides[axis]];
+    } else {
+        out[i] = rhs[rhs_i + (i_along_axis - lhs_dims[axis]) * rhs_strides[axis]];
+    }
+}
+
+extern \"C\" __global__ void bwd(
+    const size_t *info, // numel, num_dims, axis, a_dims, a_strides, b_dims, b_strides
+    $Ty *grad_lhs,
+    $Ty *grad_rhs,
+    const $Ty *grad_out
+) {
+    const size_t numel = info[0];
+    const size_t num_dims = info[1];
+    const size_t axis = info[2];
+    const size_t *lhs_dims = info + 3;
+    const size_t *lhs_strides = info + 3 + num_dims;
+    const size_t *rhs_dims = info + 3 + 2 * num_dims;
+    const size_t *rhs_strides = info + 3 + 3 * num_dims;
+
+    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i >= numel) { return; }    
+
+    // out_dims will be (..., lhs_dims[ax] + rhs_dims[ax], ...)
+
+    // striding lhs & rhs up to the concat'd axis
+    size_t i_tmp = i;
+    size_t lhs_i = 0;
+    size_t rhs_i = 0;
+    for (int d = num_dims - 1; d > axis; d--) {
+        size_t dim_i = i_tmp % lhs_dims[d];
+        lhs_i += dim_i * lhs_strides[d];
+        rhs_i += dim_i * rhs_strides[d];
+        i_tmp /= lhs_dims[d];
+    }
+
+    // figure out if we are using lhs or rhs for this `i`
+    size_t i_along_axis = i_tmp % (lhs_dims[axis] + rhs_dims[axis]);
+    i_tmp /= (lhs_dims[axis] + rhs_dims[axis]);
+
+    // striding lhs & rhs along the rest of the axes
+    for (int d = axis - 1; d >= 0;d--) {
+        size_t dim_i = i_tmp % lhs_dims[d];
+        lhs_i += dim_i * lhs_strides[d];
+        rhs_i += dim_i * rhs_strides[d];
+        i_tmp /= lhs_dims[d];
+    }
+
+    if (i_along_axis < lhs_dims[axis]) {
+        atomicAdd(grad_lhs + lhs_i + i_along_axis * lhs_strides[axis], grad_out[i]);
+    } else {
+        atomicAdd(grad_rhs + rhs_i + (i_along_axis - lhs_dims[axis]) * rhs_strides[axis], grad_out[i]);
+    }
+}
+";
diff --git a/src/tensor_ops/concat_along/mod.rs b/src/tensor_ops/concat_along/mod.rs
new file mode 100644
index 000000000..861124363
--- /dev/null
+++ b/src/tensor_ops/concat_along/mod.rs
@@ -0,0 +1,307 @@
+use crate::{shapes::*, tensor::*};
+
+mod cpu_kernel;
+#[cfg(feature = "cuda")]
+mod cuda_kernel;
+
+/// Concatenate two tensors along a given axis.
+///
+/// **Pytorch equivalent** `torch.concat`.
+///
+/// # [Const] dims **requires nightly**
+///
+/// Along Axis 0:
+/// ```ignore
+/// # use dfdx::prelude::*;
+/// # let dev: Cpu = Default::default();
+/// let a: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
+/// let b: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
+/// let _: Tensor<Rank2<6, 4>, f32, _> = (a, b).concat_along(Axis::<0>);
+/// ```
+///
+/// Along Axis 1:
+/// ```ignore
+/// # use dfdx::prelude::*;
+/// # let dev: Cpu = Default::default();
+/// let a: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
+/// let b: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
+/// let _: Tensor<Rank2<3, 8>, f32, _> = (a, b).concat_along(Axis::<1>);
+/// ```
+///
+/// # [usize] dims
+/// Along Axis 0:
+/// ```rust
+/// # use dfdx::prelude::*;
+/// # let dev: Cpu = Default::default();
+/// let a: Tensor<(usize, Const<3>), f32, _> = dev.zeros_like(&(2, Const));
+/// let b: Tensor<(usize, Const<3>), f32, _> = dev.zeros_like(&(4, Const));
+/// let _: Tensor<Rank2<6, 3>, f32, _> = (a, b).concat_along(Axis::<0>).realize().unwrap();
+/// ```
+///
+/// Along Axis 1:
+/// ```rust
+/// # use dfdx::prelude::*;
+/// # let dev: Cpu = Default::default();
+/// let a: Tensor<(Const<2>, usize), f32, _> = dev.zeros_like(&(Const, 2));
+/// let b: Tensor<(Const<2>, usize), f32, _> = dev.zeros_like(&(Const, 4));
+/// let _: Tensor<Rank2<2, 6>, f32, _> = (a, b).concat_along(Axis::<1>).realize().unwrap();
+/// ```
+pub trait TryConcatAlong<Ax>: Sized {
+    type Output;
+    type Error: std::fmt::Debug;
+    /// Concatenates self along the given axis.
+    fn concat_along(self, ax: Ax) -> Self::Output {
+        self.try_concat_along(ax).unwrap()
+    }
+    /// Fallibly concatenates self along the given axis.
+    fn try_concat_along(self, ax: Ax) -> Result<Self::Output, Self::Error>;
+}
+
+pub trait ConcatAlongKernel<E: Dtype>: DeviceStorage {
+    fn forward<A: Shape, B: Shape, C: Shape>(
+        &self,
+        ax: usize,
+        a: &Tensor<A, E, Self>,
+        b: &Tensor<B, E, Self>,
+        c: &mut Tensor<C, E, Self>,
+    ) -> Result<(), Self::Err>;
+
+    fn backward<A: Shape, B: Shape>(
+        &self,
+        ax: usize,
+        a: &GhostTensor<A, E, Self>,
+        grad_a: &mut Self::Vec<E>,
+        b: &GhostTensor<B, E, Self>,
+        grad_b: &mut Self::Vec<E>,
+        grad_out: &Self::Vec<E>,
+    ) -> Result<(), Self::Err>;
+}
+
+impl<A, B, Ax, E: Dtype, D, T: Tape<E, D>, R: Tape<E, D>> TryConcatAlong<Ax>
+    for (Tensor<A, E, D, T>, Tensor<B, E, D, R>)
+where
+    Ax: Axes<Array = [isize; 1]>,
+    D: ConcatAlongKernel<E> + ZerosTensor<E>,
+    A: Shape + HasAxes<Ax>,
+    B: Shape<Concrete = A::Concrete> + HasAxes<Ax>,
+    (A, B): TryConcatAlong<Ax>,
+    <(A, B) as TryConcatAlong<Ax>>::Output: Shape,
+    T: Merge<R>,
+{
+    type Output = Tensor<<(A, B) as TryConcatAlong<Ax>>::Output, E, D, T>;
+    type Error = D::Err;
+    fn try_concat_along(self, ax: Ax) -> Result<Self::Output, Self::Error> {
+        let (lhs, rhs) = self;
+
+        let out_shape = (*lhs.shape(), *rhs.shape()).concat_along(ax);
+        let ax = Ax::as_array()[0] as usize;
+
+        let (lhs, tape) = lhs.split_tape();
+        let (rhs, rtape) = rhs.split_tape();
+        let mut tape = tape.merge(rtape);
+
+        let mut out = lhs.device.try_zeros_like(&out_shape)?;
+        lhs.device.forward(ax, &lhs, &rhs, &mut out)?;
+
+        let lhs_ghost = lhs.ghost();
+        let rhs_ghost = rhs.ghost();
+        let out_ghost = out.ghost();
+        tape.add_backward_op(move |grads| {
+            grads.try_alloc_for(&lhs_ghost)?;
+            grads.try_alloc_for(&rhs_ghost)?;
+            grads.try_alloc_for(&out_ghost)?;
+            let (lhs_grad, rhs_grad, out_grad) =
+                grads.muts_and_ref(&lhs_ghost, &rhs_ghost, &out_ghost);
+            lhs.device
+                .backward(ax, &lhs_ghost, lhs_grad, &rhs_ghost, rhs_grad, out_grad)
+        });
+        Ok(out.put_tape(tape))
+    }
+}
+
+macro_rules! impl_concat {
+    ($Ax:expr, $NumDims:expr, [$($Head:tt),*], [$($Tail:tt),*]) => {
+        impl<A: Dim, B: Dim, $($Head: Dim, )* $($Tail: Dim, )*> TryConcatAlong<Axis<$Ax>>
+            for (
+                ($($Head, )* A, $($Tail, )*),
+                ($($Head, )* B, $($Tail, )*),
+            )
+        where
+            A: std::ops::Add<B>,
+            <A as std::ops::Add<B>>::Output: Dim,
+            {
+                type Output = (
+                    $($Head, )*
+                    <A as std::ops::Add<B>>::Output,
+                    $($Tail, )*
+                );
+                type Error = std::convert::Infallible;
+                fn try_concat_along(self, _: Axis<$Ax>) -> Result<Self::Output, Self::Error> {
+                    let (lhs, rhs) = self;
+                    let lhs_dims = lhs.concrete();
+                    let rhs_dims = rhs.concrete();
+                    for i in 0..$NumDims {
+                        if i != $Ax {
+                            assert_eq!(lhs_dims[i], rhs_dims[i]);
+                        }
+                    }
+                    let mut out_dims = lhs_dims;
+                    out_dims[$Ax] += rhs_dims[$Ax];
+                    Ok(Self::Output::from_concrete(&out_dims).unwrap())
+                }
+            }
+    };
+}
+
+impl_concat!(0, 1, [], []);
+impl_concat!(0, 2, [], [D1]);
+impl_concat!(0, 3, [], [D1, D2]);
+impl_concat!(0, 4, [], [D1, D2, D3]);
+impl_concat!(0, 5, [], [D1, D2, D3, D4]);
+impl_concat!(0, 6, [], [D1, D2, D3, D4, D5]);
+
+impl_concat!(1, 2, [D0], []);
+impl_concat!(1, 3, [D0], [D2]);
+impl_concat!(1, 4, [D0], [D2, D3]);
+impl_concat!(1, 5, [D0], [D2, D3, D4]);
+impl_concat!(1, 6, [D0], [D2, D3, D4, D5]);
+
+impl_concat!(2, 3, [D0, D1], []);
+impl_concat!(2, 4, [D0, D1], [D3]);
+impl_concat!(2, 5, [D0, D1], [D3, D4]);
+impl_concat!(2, 6, [D0, D1], [D3, D4, D5]);
+
+impl_concat!(3, 4, [D0, D1, D2], []);
+impl_concat!(3, 5, [D0, D1, D2], [D4]);
+impl_concat!(3, 6, [D0, D1, D2], [D4, D5]);
+
+impl_concat!(4, 5, [D0, D1, D2, D3], []);
+impl_concat!(4, 6, [D0, D1, D2, D3], [D5]);
+
+impl_concat!(5, 6, [D0, D1, D2, D3, D4], []);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::{tensor_ops::*, tests::*};
+
+    #[test]
+    fn test_concat_ax_0() {
+        let dev: TestDevice = Default::default();
+        let a: Tensor<Rank3<2, 3, 4>, TestDtype, _> = dev.sample_normal();
+        let b: Tensor<Rank3<3, 3, 4>, TestDtype, _> = dev.sample_normal();
+        let a_dyn = a
+            .leaky_trace()
+            .realize::<(usize, Const<3>, Const<4>)>()
+            .unwrap();
+        let b_dyn = b.clone().realize::<(usize, Const<3>, Const<4>)>().unwrap();
+        let c = (a_dyn, b_dyn).concat_along(Axis::<0>);
+        let c = c.realize::<(Const<5>, Const<3>, Const<4>)>().unwrap();
+        let a_arr = a.array();
+        let b_arr = b.array();
+        let c_arr = c.array();
+        println!("{a_arr:?}");
+        println!("{b_arr:?}");
+        println!("{c_arr:?}");
+        assert_eq!(c_arr[0], a_arr[0]);
+        assert_eq!(c_arr[1], a_arr[1]);
+        assert_eq!(c_arr[2], b_arr[0]);
+        assert_eq!(c_arr[3], b_arr[1]);
+        assert_eq!(c_arr[4], b_arr[2]);
+        let concat_grads = c.exp().sum().backward();
+        let a_grads = a.leaky_trace().exp().sum().backward();
+        let b_grads = b.leaky_trace().exp().sum().backward();
+        assert_close_to_tensor!(concat_grads.get(&a), a_grads.get(&a));
+        assert_close_to_tensor!(concat_grads.get(&b), b_grads.get(&b));
+    }
+
+    #[test]
+    fn test_concat_ax_1() {
+        let dev: TestDevice = Default::default();
+        let a: Tensor<Rank3<2, 2, 4>, TestDtype, _> = dev.sample_normal();
+        let b: Tensor<Rank3<2, 3, 4>, TestDtype, _> = dev.sample_normal();
+        let a_dyn = a
+            .leaky_trace()
+            .realize::<(Const<2>, usize, Const<4>)>()
+            .unwrap();
+        let b_dyn = b.clone().realize::<(Const<2>, usize, Const<4>)>().unwrap();
+        let c = (a_dyn, b_dyn).concat_along(Axis::<1>);
+        let c = c.realize::<(Const<2>, Const<5>, Const<4>)>().unwrap();
+        let a_arr = a.array();
+        let b_arr = b.array();
+        let c_arr = c.array();
+        for i in 0..2 {
+            assert_eq!(c_arr[i][0], a_arr[i][0]);
+            assert_eq!(c_arr[i][1], a_arr[i][1]);
+            assert_eq!(c_arr[i][2], b_arr[i][0]);
+            assert_eq!(c_arr[i][3], b_arr[i][1]);
+            assert_eq!(c_arr[i][4], b_arr[i][2]);
+        }
+        let concat_grads = c.exp().sum().backward();
+        let a_grads = a.leaky_trace().exp().sum().backward();
+        let b_grads = b.leaky_trace().exp().sum().backward();
+        assert_close_to_tensor!(concat_grads.get(&a), a_grads.get(&a));
+        assert_close_to_tensor!(concat_grads.get(&b), b_grads.get(&b));
+    }
+
+    #[test]
+    fn test_concat_ax_2() {
+        let dev: TestDevice = Default::default();
+        let a: Tensor<Rank3<2, 3, 2>, TestDtype, _> = dev.sample_normal();
+        let b: Tensor<Rank3<2, 3, 3>, TestDtype, _> = dev.sample_normal();
+        let a_dyn = a
+            .leaky_trace()
+            .realize::<(Const<2>, Const<3>, usize)>()
+            .unwrap();
+        let b_dyn = b.clone().realize::<(Const<2>, Const<3>, usize)>().unwrap();
+        let c = (a_dyn, b_dyn).concat_along(Axis::<2>);
+        let c = c.realize::<(Const<2>, Const<3>, Const<5>)>().unwrap();
+        let a_arr = a.array();
+        let b_arr = b.array();
+        let c_arr = c.array();
+        for i in 0..2 {
+            for j in 0..3 {
+                assert_eq!(c_arr[i][j][0], a_arr[i][j][0]);
+                assert_eq!(c_arr[i][j][1], a_arr[i][j][1]);
+                assert_eq!(c_arr[i][j][2], b_arr[i][j][0]);
+                assert_eq!(c_arr[i][j][3], b_arr[i][j][1]);
+                assert_eq!(c_arr[i][j][4], b_arr[i][j][2]);
+            }
+        }
+        let concat_grads = c.exp().sum().backward();
+        let a_grads = a.leaky_trace().exp().sum().backward();
+        let b_grads = b.leaky_trace().exp().sum().backward();
+        assert_close_to_tensor!(concat_grads.get(&a), a_grads.get(&a));
+        assert_close_to_tensor!(concat_grads.get(&b), b_grads.get(&b));
+    }
+
+    #[test]
+    fn test_concat_shape() {
+        let a: (usize, Const<5>) = (5, Const);
+        let b: (usize, Const<5>) = (3, Const);
+        assert_eq!((a, b).concat_along(Axis::<0>), (8, Const::<5>));
+
+        let a: (Const<5>, Const<5>) = (Const, Const);
+        let b: (usize, Const<5>) = (3, Const);
+        assert_eq!((a, b).concat_along(Axis::<0>), (8, Const::<5>));
+
+        let a: (usize, Const<5>) = (5, Const);
+        let b: (Const<3>, Const<5>) = (Const, Const);
+        assert_eq!((a, b).concat_along(Axis::<0>), (8, Const::<5>));
+
+        #[cfg(feature = "nightly")]
+        {
+            let a: (Const<5>, Const<5>) = (Const, Const);
+            let b: (Const<3>, Const<5>) = (Const, Const);
+            assert_eq!((a, b).concat_along(Axis::<0>), (Const::<8>, Const::<5>));
+        }
+    }
+
+    #[test]
+    #[should_panic = "left: `10`,\n right: `7`"]
+    fn test_concat_shape_fails() {
+        let a = (5, 10);
+        let b = (3, 7);
+        (a, b).concat_along(Axis::<0>);
+    }
+}
diff --git a/src/tensor_ops/mod.rs b/src/tensor_ops/mod.rs
index aaae02be1..9dfd8b45a 100644
--- a/src/tensor_ops/mod.rs
+++ b/src/tensor_ops/mod.rs
@@ -160,6 +160,7 @@ mod choose;
 mod clamp;
 mod cmp;
 mod concat;
+mod concat_along;
 mod cos;
 mod div;
 mod dropout;
@@ -213,7 +214,9 @@ pub use broadcast_to::BroadcastTo;
 pub use choose::ChooseFrom;
 pub use clamp::clamp;
 pub use cmp::{eq, ge, gt, le, lt, ne, TryEq, TryGe, TryGt, TryLe, TryLt, TryNe};
+#[allow(deprecated)]
 pub use concat::TryConcat;
+pub use concat_along::TryConcatAlong;
 pub use cos::cos;
 pub use div::{div, TryDiv};
 pub use dropout::dropout;
diff --git a/src/tensor_ops/utilities/device.rs b/src/tensor_ops/utilities/device.rs
index e20cee99d..3cb001ac1 100644
--- a/src/tensor_ops/utilities/device.rs
+++ b/src/tensor_ops/utilities/device.rs
@@ -14,6 +14,7 @@ pub trait Device<E: Dtype>:
     // appends
     + super::super::stack::StackKernel<E>
     + super::super::concat::ConcatKernel<E>
+    + super::super::concat_along::ConcatAlongKernel<E>
 
     // optimizers
     + crate::optim::AdamKernel<E>


Test patch: