zyx 0.17.0

Zyx machine learning library
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only WITH Classpath-exception-2.0

//! Few traits that describe shapes, axes, padding, etc.

use crate::{error::ZyxError, tensor::Axis};

/// Type alias for dimension values (i64)
///
/// # Convention: `-1` means dynamic/symbolic
///
/// A `Dim` of `-1` does NOT mean an invalid dimension — it marks a **dynamic
/// (symbolic) dimension** whose length is only known at kernel launch time
/// (supplied as a `Param { kind: Variable }` scalar). Any nonnegative value
/// is a static, known length (emitted as a `Const`).
pub type Dim = i64;
/// Type alias for axis indices (usize)
pub type UAxis = usize;

pub fn into_axis(axis: Axis, rank: UAxis) -> Result<UAxis, ZyxError> {
    TryInto::<Axis>::try_into(rank).map_or_else(
        |_| Err(ZyxError::ShapeError(format!("Axis {axis} is out of range of rank {rank}").into())),
        |rank2| {
            TryInto::<UAxis>::try_into(axis + rank2).map_or_else(
                |_| Err(ZyxError::ShapeError(format!("Axis {axis} is out of range of rank {rank}").into())),
                |a| {
                    if a < 2 * rank {
                        Ok(a % rank)
                    } else {
                        Err(ZyxError::ShapeError(format!("Axis {axis} is out of range of rank {rank}").into()))
                    }
                },
            )
        },
    )
}

pub fn into_axes(axes: impl IntoIterator<Item = Axis>, rank: UAxis) -> Result<Vec<UAxis>, ZyxError> {
    let mut res = Vec::with_capacity(rank);
    let mut visited = std::collections::BTreeSet::new();
    for axis in axes {
        let a = into_axis(axis, rank)?;
        if visited.insert(a) {
            res.push(a);
        }
    }
    if res.is_empty() {
        return Ok((0..rank).collect());
    }
    Ok(res)
}

#[must_use]
pub fn permute<T: Clone>(shape: &[T], axes: &[UAxis]) -> Vec<T> {
    debug_assert_eq!(shape.len(), axes.len());
    axes.iter().map(|a| shape[*a].clone()).collect()
}

pub fn pad(shape: &mut [Dim], padding: &[(i64, i64)]) {
    let mut i = 0;
    for d in shape.iter_mut() {
        *d = Dim::try_from(i64::try_from(*d).unwrap() + padding[i].0 + padding[i].1).unwrap();
        i += 1;
        if i >= padding.len() {
            break;
        }
    }
}