embedded_nn/
transpose.rs

1//! Tensor transpose operations.
2
3use crate::types::{Dims, Result};
4
5/// Transposes a 2D matrix (rows x cols) for int8 tensors.
6pub fn transpose_2d_s8(rows: usize, cols: usize, input: &[i8], output: &mut [i8]) -> Result<()> {
7    for r in 0..rows {
8        for c in 0..cols {
9            output[c * rows + r] = input[r * cols + c];
10        }
11    }
12    Ok(())
13}
14
15/// Transposes a 4D tensor (N, H, W, C) -> (N, W, H, C).
16pub fn transpose_spatial_s8(input_dims: &Dims, input: &[i8], output: &mut [i8]) -> Result<()> {
17    let batches = input_dims.n as usize;
18    let input_h = input_dims.h as usize;
19    let input_w = input_dims.w as usize;
20    let channels = input_dims.c as usize;
21
22    for b in 0..batches {
23        for y in 0..input_h {
24            for x in 0..input_w {
25                let in_idx = ((b * input_h + y) * input_w + x) * channels;
26                let out_idx = ((b * input_w + x) * input_h + y) * channels;
27
28                output[out_idx..out_idx + channels]
29                    .copy_from_slice(&input[in_idx..in_idx + channels]);
30            }
31        }
32    }
33    Ok(())
34}
35
36#[cfg(test)]
37mod tests {
38    use super::*;
39
40    #[test]
41    fn test_transpose_2d_s8() {
42        let input = [1i8, 2i8, 3i8, 4i8, 5i8, 6i8]; // 2 rows x 3 cols
43        let mut out = [0i8; 6]; // 3 rows x 2 cols
44        transpose_2d_s8(2, 3, &input, &mut out).unwrap();
45        assert_eq!(out, [1, 4, 2, 5, 3, 6]);
46    }
47}
embedded_nn/transpose.rs

embedded_nn/
transpose.rs
