vision_rs/models/yolo/yolo26/blocks/
concat.rs

1/*
2 * Copyright 2026 Teenygrad
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17
18use teeny_core::graph::{Op, SymTensor};
19
20/// Channel-wise concatenation of N NCHW tensors.
21///
22/// All inputs must share the same B, H, W. Channels are concatenated in order.
23/// Output shape: `[B, C0+C1+...+CN, H, W]`.
24pub fn concat() -> impl Fn(Vec<SymTensor>) -> SymTensor {
25    |tensors| {
26        assert!(!tensors.is_empty(), "concat requires at least one input");
27        let c_total: usize = tensors.iter()
28            .map(|t| t.shape[1].expect("channel dim must be known"))
29            .sum();
30        let first = &tensors[0];
31        let shape = vec![first.shape[0], Some(c_total), first.shape[2], first.shape[3]];
32        let inputs: Vec<usize> = tensors.iter().map(|t| t.node_id).collect();
33        let node_id = first.graph.borrow_mut().add_node(
34            Op::ChannelCat { c_total }, inputs, first.dtype, shape.clone(),
35        );
36        SymTensor { node_id, graph: first.graph.clone(), dtype: first.dtype, shape }
37    }
38}
vision_rs/models/yolo/yolo26/blocks/concat.rs

vision_rs/models/yolo/yolo26/blocks/
concat.rs
