maidenx_tensor 0.1.5

maidenx tensor
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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use crate::Tensor;
use maidenx_core::{
    error::{Error, Result},
    scalar::Scalar,
};

pub fn select_dim(src: &Tensor, dim: impl Into<Scalar>, index: impl Into<Scalar>) -> Result<Tensor> {
    let dim_i32 = dim.into().as_i32();
    let index_i32 = index.into().as_i32();

    let dim: usize = if dim_i32 < 0 {
        (src.ndim() as i32 + dim_i32) as usize
    } else {
        dim_i32 as usize
    };

    if dim >= src.ndim() {
        return Err(Error::DimensionOutOfBounds {
            dim: dim as i32,
            ndim: src.ndim(),
        });
    }

    let start = index_i32;
    let end = Some(index_i32 + 1);
    let step = 1;

    let sliced = src.slice(dim, start, end, step)?;
    sliced.squeeze(dim)
}

pub fn get_index(src: &Tensor, indices: &[usize]) -> Result<Scalar> {
    if indices.len() != src.ndim() {
        return Err(Error::InvalidArgument(format!(
            "Indices length ({}) must match tensor dimensions ({})",
            indices.len(),
            src.ndim()
        )));
    }

    let mut usize_indices = Vec::with_capacity(indices.len());
    for (dim, &idx) in indices.iter().enumerate() {
        let idx_i32 = idx as i32;
        let dim_size = src.dim_size(dim).unwrap_or(0) as i32;

        let actual_idx = if idx_i32 < 0 { dim_size + idx_i32 } else { idx_i32 };

        if actual_idx < 0 || actual_idx >= dim_size {
            return Err(Error::IndexOutOfBounds {
                index: actual_idx as usize,
                size: dim_size as usize,
            });
        }

        usize_indices.push(actual_idx as usize);
    }

    let mut flat_idx = 0;
    for (dim, &idx) in usize_indices.iter().enumerate() {
        flat_idx += idx * src.strides()[dim];
    }

    flat_idx += src.offset();
    src.buffer().read_scalar(flat_idx)
}

pub fn set_index(src: &mut Tensor, indices: &[usize], data: impl Into<Scalar>) -> Result<()> {
    if indices.len() != src.ndim() {
        return Err(Error::InvalidArgument(format!(
            "Indices length ({}) must match tensor dimensions ({})",
            indices.len(),
            src.ndim()
        )));
    }

    for (dim, &idx) in indices.iter().enumerate() {
        if idx >= src.shape()[dim] {
            return Err(Error::IndexOutOfBounds {
                index: idx,
                size: src.shape()[dim],
            });
        }
    }

    let mut flat_idx = 0;
    for (dim, &idx) in indices.iter().enumerate() {
        flat_idx += idx * src.strides()[dim];
    }

    flat_idx += src.offset();

    let scalar = data.into();
    src.buffer_mut().write_scalar(flat_idx, scalar)
}

pub fn add_at_index(src: &mut Tensor, indices: &[usize], data: impl Into<Scalar>) -> Result<()> {
    if indices.len() != src.ndim() {
        return Err(Error::InvalidArgument(format!(
            "Indices length ({}) must match tensor dimensions ({})",
            indices.len(),
            src.ndim()
        )));
    }

    for (dim, &idx) in indices.iter().enumerate() {
        if idx >= src.shape()[dim] {
            return Err(Error::IndexOutOfBounds {
                index: idx,
                size: src.shape()[dim],
            });
        }
    }

    let mut flat_idx = 0;
    for (dim, &idx) in indices.iter().enumerate() {
        flat_idx += idx * src.strides()[dim];
    }

    flat_idx += src.offset();

    let add_value = data.into();
    let dtype = src.dtype();

    let current_value = src.buffer().read_scalar(flat_idx)?;

    let new_value = match dtype {
        maidenx_core::dtype::DType::BF16 => Scalar::from(current_value.as_bf16() + add_value.as_bf16()),
        maidenx_core::dtype::DType::F16 => Scalar::from(current_value.as_f16() + add_value.as_f16()),
        maidenx_core::dtype::DType::F32 => Scalar::from(current_value.as_f32() + add_value.as_f32()),
        maidenx_core::dtype::DType::F64 => Scalar::from(current_value.as_f64() + add_value.as_f64()),
        maidenx_core::dtype::DType::BOOL => Scalar::from(current_value.as_bool() || add_value.as_bool()),
        maidenx_core::dtype::DType::U8 => Scalar::from(current_value.as_u8() + add_value.as_u8()),
        maidenx_core::dtype::DType::U16 => Scalar::from(current_value.as_u16() + add_value.as_u16()),
        maidenx_core::dtype::DType::U32 => Scalar::from(current_value.as_u32() + add_value.as_u32()),
        maidenx_core::dtype::DType::U64 => Scalar::from(current_value.as_u64() + add_value.as_u64()),
        maidenx_core::dtype::DType::I8 => Scalar::from(current_value.as_i8() + add_value.as_i8()),
        maidenx_core::dtype::DType::I16 => Scalar::from(current_value.as_i16() + add_value.as_i16()),
        maidenx_core::dtype::DType::I32 => Scalar::from(current_value.as_i32() + add_value.as_i32()),
        maidenx_core::dtype::DType::I64 => Scalar::from(current_value.as_i64() + add_value.as_i64()),
    };

    src.buffer_mut().write_scalar(flat_idx, new_value)
}