coreylowman · coreylowman · Mar 26, 2023 · Mar 26, 2023 · Mar 26, 2023
diff --git a/src/tensor_ops/concat/mod.rs b/src/tensor_ops/concat/mod.rs
@@ -102,19 +102,19 @@ pub trait ConcatShape<Rhs: Shape>: Shape {
 
 macro_rules! impl_concat {
     ([$($Dims:tt $Idx:tt),*]) => {
-impl<A: Dim, B: Dim, $($Dims: Dim, )*> ConcatShape<(A, $($Dims, )*)>
-    for (B, $($Dims, )*)
-where
-    A: std::ops::Add<B>,
-    <A as std::ops::Add<B>>::Output: Dim,
-{
-    type Catted = (<A as std::ops::Add<B>>::Output, $($Dims, )*);
+        impl<A: Dim, B: Dim, $($Dims: Dim, )*> ConcatShape<(A, $($Dims, )*)>
+            for (B, $($Dims, )*)
+        where
+            A: std::ops::Add<B>,
+            <A as std::ops::Add<B>>::Output: Dim,
+        {
+            type Catted = (<A as std::ops::Add<B>>::Output, $($Dims, )*);
 
-    fn concat_shape(&self, rhs: &(A, $($Dims, )*)) -> Self::Catted {
-        $(assert_eq!(self.$Idx, rhs.$Idx);)*
-        (rhs.0 + self.0, $(self.$Idx, )*)
-    }
-}
+            fn concat_shape(&self, rhs: &(A, $($Dims, )*)) -> Self::Catted {
+                $(assert_eq!(self.$Idx, rhs.$Idx);)*
+                (rhs.0 + self.0, $(self.$Idx, )*)
+            }
+        }
     };
 }
 
@@ -125,6 +125,20 @@ impl_concat!([D1 1, D2 2, D3 3]);
 impl_concat!([D1 1, D2 2, D3 3, D4 4]);
 impl_concat!([D1 1, D2 2, D3 3, D4 4, D5 5]);
 
+impl<const N: usize> ConcatShape<[usize; N]> for [usize; N]
+where
+    [usize; N]: Shape,
+{
+    type Catted = [usize; N];
+
+    fn concat_shape(&self, rhs: &[usize; N]) -> [usize; N] {
+        assert_eq!(self[1..], rhs[1..]);
+        let mut out = *self;
+        out[0] = self[0] + rhs[0];
+        out
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

diff --git a/src/tensor_ops/utilities/device.rs b/src/tensor_ops/utilities/device.rs
@@ -11,7 +11,9 @@ pub trait Device<E: Dtype>:
     + crate::tensor::TensorFromVec<E>
     + crate::tensor::TensorFromVec<usize>
 
-    + crate::tensor_ops::stack::StackKernel<E>
+    // appends
+    + super::super::stack::StackKernel<E>
+    + super::super::concat::ConcatKernel<E>
 
     // optimizers
     + crate::optim::AdamKernel<E>