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
80            | DType::Q4_1
81            | DType::Q5_0
82            | DType::Q5_1
83            | DType::Q8_0
84            | DType::Q2_K
85            | DType::Q3_K
86            | DType::Q4_K
87            | DType::Q5_K
88            | DType::Q6_K
89            | DType::Q8_K => {
90                unreachable!("quantized dtypes are rejected above")
91            }
92        };
93
94        Ok(Tensor {
95            storage: std::sync::Arc::new(storage),
96            strides: out_shape.strides(),
97            shape: out_shape,
98            offset: 0,
99            device: first.device,
100        })
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn concat_along_the_last_axis_interleaves_rows_correctly() {
110        // [[1,2],[3,4]] concat [[5,6],[7,8]] along axis 1 -> [[1,2,5,6],[3,4,7,8]]
111        let a = Tensor::from_f32(vec![1.0, 2.0, 3.0, 4.0], [2, 2]).unwrap();
112        let b = Tensor::from_f32(vec![5.0, 6.0, 7.0, 8.0], [2, 2]).unwrap();
113        let c = Tensor::concat(&[a, b], 1).unwrap();
114        assert_eq!(c.shape().dims(), &[2, 4]);
115        assert_eq!(c.to_vec_f32().unwrap(), vec![1.0, 2.0, 5.0, 6.0, 3.0, 4.0, 7.0, 8.0]);
116    }
117
118    #[test]
119    fn concat_along_the_first_axis_stacks_rows() {
120        let a = Tensor::from_f32(vec![1.0, 2.0], [1, 2]).unwrap();
121        let b = Tensor::from_f32(vec![3.0, 4.0], [1, 2]).unwrap();
122        let c = Tensor::concat(&[a, b], 0).unwrap();
123        assert_eq!(c.shape().dims(), &[2, 2]);
124        assert_eq!(c.to_vec_f32().unwrap(), vec![1.0, 2.0, 3.0, 4.0]);
125    }
126
127    #[test]
128    fn concat_three_tensors_matches_hand_computation() {
129        let a = Tensor::from_f32(vec![1.0], [1]).unwrap();
130        let b = Tensor::from_f32(vec![2.0], [1]).unwrap();
131        let c = Tensor::from_f32(vec![3.0], [1]).unwrap();
132        let out = Tensor::concat(&[a, b, c], 0).unwrap();
133        assert_eq!(out.to_vec_f32().unwrap(), vec![1.0, 2.0, 3.0]);
134    }
135
136    #[test]
137    fn concat_rejects_mismatched_non_concat_dimensions() {
138        let a = Tensor::from_f32(vec![1.0, 2.0], [1, 2]).unwrap();
139        let b = Tensor::from_f32(vec![1.0, 2.0, 3.0], [1, 3]).unwrap();
140        assert!(matches!(Tensor::concat(&[a, b], 0), Err(Error::ShapeMismatch { .. })));
141    }
142
143    #[test]
144    fn concat_of_non_contiguous_views_still_produces_correct_data() {
145        let base = Tensor::from_f32(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [2, 3]).unwrap();
146        let transposed = base.transpose(0, 1).unwrap(); // [3, 2]: [[1,4],[2,5],[3,6]]
147        let other = Tensor::from_f32(vec![9.0, 9.0], [1, 2]).unwrap();
148        let out = Tensor::concat(&[transposed, other], 0).unwrap();
149        assert_eq!(out.to_vec_f32().unwrap(), vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0, 9.0, 9.0]);
150    }
151}
kopitiam_tensor/tensor/concat.rs

kopitiam_tensor/tensor/
concat.rs
