kopitiam_tensor/tensor/
concat.rs

1//! Concatenating tensors along an axis.
2//!
3//! Unlike `narrow`/`transpose`/`broadcast_to`, concatenation cannot be
4//! zero-copy — the output genuinely interleaves bytes from multiple
5//! inputs — so this is the one shape op in the `tensor` module family that
6//! always allocates.
7
8use kopitiam_core::{DType, Error, Result, Shape};
9
10use crate::storage::Storage;
11
12use super::Tensor;
13
14impl Tensor {
15    /// Concatenates `tensors` along `dim`. All tensors must share a dtype
16    /// (any non-quantized dtype — this is pure data movement, not
17    /// arithmetic, so it is not restricted to `f32` the way the math ops
18    /// are) and must agree on every dimension except `dim`.
19    pub fn concat(tensors: &[Tensor], dim: usize) -> Result<Tensor> {
20        let Some(first) = tensors.first() else {
21            // No tensor to take a dtype/rank from; there is no sensible
22            // shape for an empty concatenation.
23            return Err(Error::ShapeMismatch { expected: Shape::scalar(), actual: Shape::scalar() });
24        };
25        let dtype = first.dtype();
26        if dtype.is_quantized() {
27            return Err(Error::QuantizedElementAccess { dtype, block_size: dtype.block_size() });
28        }
29        let rank = first.rank();
30        if dim >= rank {
31            return Err(Error::IndexOutOfBounds { dim, index: dim, len: rank });
32        }
33        for t in tensors {
34            if t.dtype() != dtype {
35                return Err(Error::DTypeMismatch { expected: dtype, actual: t.dtype() });
36            }
37            if t.rank() != rank {
38                return Err(Error::ShapeMismatch { expected: first.shape.clone(), actual: t.shape.clone() });
39            }
40            for (d, (&a, &b)) in first.shape.dims().iter().zip(t.shape.dims()).enumerate() {
41                if d != dim && a != b {
42                    return Err(Error::ShapeMismatch { expected: first.shape.clone(), actual: t.shape.clone() });
43                }
44            }
45        }
46
47        let mut out_dims = first.shape.dims().to_vec();
48        out_dims[dim] = tensors.iter().map(|t| t.shape.dims()[dim]).sum();
49        let out_shape = Shape::new(out_dims.clone());
50        let outer: usize = out_dims[..dim].iter().product();
51        let inner: usize = out_dims[dim + 1..].iter().product();
52
53        let contiguous: Vec<Tensor> = tensors.iter().map(Tensor::contiguous).collect::<Result<_>>()?;
54
55        // One macro arm per `Storage` variant instead of five hand-written
56        // copies of the same "slice out each input's chunk along `dim`,
57        // append to `out`" loop, which differ only in element type.
58        macro_rules! concat_variant {
59            ($variant:ident, $ty:ty) => {{
60                let mut out: Vec<$ty> = Vec::with_capacity(out_shape.elem_count());
61                for o in 0..outer {
62                    for t in &contiguous {
63                        let Storage::$variant(src) = t.storage.as_ref() else { unreachable!() };
64                        let len_d = t.shape.dims()[dim];
65                        let start = o * len_d * inner;
66                        out.extend_from_slice(&src[start..start + len_d * inner]);
67                    }
68                }
69                Storage::$variant(out)
70            }};
71        }
72
73        let storage = match dtype {
74            DType::F32 => concat_variant!(F32, f32),
75            DType::F16 => concat_variant!(F16, u16),
76            DType::BF16 => concat_variant!(BF16, u16),
77            DType::I8 => concat_variant!(I8, i8),
78            DType::I32 => concat_variant!(I32, i32),
79            DType::Q4_0 | DType::Q4_1 | DType::Q5_0 | DType::Q5_1 | DType::Q8_0 => {
80                unreachable!("quantized dtypes are rejected above")
81            }
82        };
83
84        Ok(Tensor {
85            storage: std::sync::Arc::new(storage),
86            strides: out_shape.strides(),
87            shape: out_shape,
88            offset: 0,
89            device: first.device,
90        })
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97
98    #[test]
99    fn concat_along_the_last_axis_interleaves_rows_correctly() {
100        // [[1,2],[3,4]] concat [[5,6],[7,8]] along axis 1 -> [[1,2,5,6],[3,4,7,8]]
101        let a = Tensor::from_f32(vec![1.0, 2.0, 3.0, 4.0], [2, 2]).unwrap();
102        let b = Tensor::from_f32(vec![5.0, 6.0, 7.0, 8.0], [2, 2]).unwrap();
103        let c = Tensor::concat(&[a, b], 1).unwrap();
104        assert_eq!(c.shape().dims(), &[2, 4]);
105        assert_eq!(c.to_vec_f32().unwrap(), vec![1.0, 2.0, 5.0, 6.0, 3.0, 4.0, 7.0, 8.0]);
106    }
107
108    #[test]
109    fn concat_along_the_first_axis_stacks_rows() {
110        let a = Tensor::from_f32(vec![1.0, 2.0], [1, 2]).unwrap();
111        let b = Tensor::from_f32(vec![3.0, 4.0], [1, 2]).unwrap();
112        let c = Tensor::concat(&[a, b], 0).unwrap();
113        assert_eq!(c.shape().dims(), &[2, 2]);
114        assert_eq!(c.to_vec_f32().unwrap(), vec![1.0, 2.0, 3.0, 4.0]);
115    }
116
117    #[test]
118    fn concat_three_tensors_matches_hand_computation() {
119        let a = Tensor::from_f32(vec![1.0], [1]).unwrap();
120        let b = Tensor::from_f32(vec![2.0], [1]).unwrap();
121        let c = Tensor::from_f32(vec![3.0], [1]).unwrap();
122        let out = Tensor::concat(&[a, b, c], 0).unwrap();
123        assert_eq!(out.to_vec_f32().unwrap(), vec![1.0, 2.0, 3.0]);
124    }
125
126    #[test]
127    fn concat_rejects_mismatched_non_concat_dimensions() {
128        let a = Tensor::from_f32(vec![1.0, 2.0], [1, 2]).unwrap();
129        let b = Tensor::from_f32(vec![1.0, 2.0, 3.0], [1, 3]).unwrap();
130        assert!(matches!(Tensor::concat(&[a, b], 0), Err(Error::ShapeMismatch { .. })));
131    }
132
133    #[test]
134    fn concat_of_non_contiguous_views_still_produces_correct_data() {
135        let base = Tensor::from_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]).unwrap();
136        let transposed = base.transpose(0, 1).unwrap(); // [3, 2]: [[1,4],[2,5],[3,6]]
137        let other = Tensor::from_f32(vec![9.0, 9.0], [1, 2]).unwrap();
138        let out = Tensor::concat(&[transposed, other], 0).unwrap();
139        assert_eq!(out.to_vec_f32().unwrap(), vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0, 9.0, 9.0]);
140    }
141}
kopitiam_tensor/tensor/concat.rs

kopitiam_tensor/tensor/
concat.rs
