svod-tensor 0.1.0-alpha.3

High-level lazy tensor API for the Svod ML compiler
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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199

//! Distribution wrappers around `Tensor::rand`.
//!
//! | Method | Formula |
//! |---|---|
//! | `uniform(shape, low, high)` | `(high - low) * rand + low` |
//! | `randn(shape)` | Box-Muller: `cos(2π·u₁) · √(-2·ln(1 - u₂))` |
//! | `normal(shape, mean, std)` | `std · randn + mean` |
//! | `randint(shape, low, high)` | `((high-low)·rand).cast(int32) + low` |
//! | `scaled_uniform(shape)` | `uniform(-1, 1) · prod(shape)^-½` |
//! | `glorot_uniform(shape)` | `uniform(-b, b)`, `b = √(6 / (shape[0] + prod(shape[1..])))` |
//! | `kaiming_uniform(shape, a)` | `uniform(-b, b)`, `b = √(6 / ((1+a²) · prod(shape[1..])))` |
//! | `kaiming_normal(shape, a)` | `randn · √(2 / ((1+a²) · prod(shape[1..])))` |

use svod_dtype::DType;
use svod_ir::ConstValue;

use crate::{Error, Result, Tensor};

const TWO_PI: f64 = 2.0 * std::f64::consts::PI;

fn fan_in(shape: &[usize]) -> usize {
    // fan_in = prod(shape[1..]). For 1D inputs (e.g., bias) the product over
    // the empty slice is 1.
    shape.iter().skip(1).copied().product::<usize>().max(1)
}

impl Tensor {
    /// Uniform `[low, high)` random tensor, float32, on the default (CPU) device.
    ///
    /// Convenience wrapper around [`Tensor::uniform_with_dtype`] with f32 output.
    pub fn uniform(shape: &[usize], low: f64, high: f64) -> Result<Tensor> {
        Self::uniform_with_dtype(shape, low, high, DType::Float32)
    }

    /// Uniform `[low, high)` random tensor with explicit float dtype.
    ///
    /// Generates a `[0, 1)` sample at f32, scales by `(high - low)`, **casts
    /// to the target dtype**, then adds `low`. Casting before the offset
    /// keeps the addition honest in low-precision targets (f16/bf16) where
    /// `low` might otherwise be lost to rounding if applied at f32.
    pub fn uniform_with_dtype(shape: &[usize], low: f64, high: f64, dtype: DType) -> Result<Tensor> {
        if low >= high {
            return Err(Error::ParamRange {
                op: "Tensor::uniform",
                param: "low/high",
                value: format!("low={low}, high={high}"),
                constraint: "low < high",
            });
        }
        let u = Tensor::rand(shape)?;
        let scale = u.broadcast_scalar(ConstValue::Float(high - low))?;
        let scaled = u.try_mul(&scale)?.cast(dtype)?;
        let offset = scaled.broadcast_scalar(ConstValue::Float(low))?;
        scaled.try_add(&offset)
    }

    /// Standard normal `N(0, 1)` random tensor (float32, Box-Muller).
    ///
    /// Each output element draws from two `[0, 1)` uniforms via one combined
    /// `rand([2, *shape])` call, so the RNG counter advances exactly once per
    /// `randn` invocation regardless of `shape`.
    pub fn randn(shape: &[usize]) -> Result<Tensor> {
        // src = rand([2, *shape])  →  one counter advance for two halves.
        let mut combined_shape: Vec<usize> = Vec::with_capacity(shape.len() + 1);
        combined_shape.push(2);
        combined_shape.extend_from_slice(shape);
        let src = Tensor::rand(&combined_shape)?;

        // u1 = src[0:1, …] reshaped to shape; u2 = src[1:2, …] reshaped to shape.
        let mut shrink_u1: Vec<Option<(isize, isize)>> = Vec::with_capacity(combined_shape.len());
        shrink_u1.push(Some((0, 1)));
        shrink_u1.extend(std::iter::repeat_n(None, shape.len()));
        let mut shrink_u2: Vec<Option<(isize, isize)>> = Vec::with_capacity(combined_shape.len());
        shrink_u2.push(Some((1, 2)));
        shrink_u2.extend(std::iter::repeat_n(None, shape.len()));
        let target_shape: Vec<isize> = shape.iter().map(|&d| d as isize).collect();
        let u1 = src.try_shrink(shrink_u1)?.try_reshape(&target_shape)?;
        let u2 = src.try_shrink(shrink_u2)?.try_reshape(&target_shape)?;

        // Box-Muller: cos(2π·u1) · √(-2·ln(1 - u2))
        let two_pi = u1.broadcast_scalar(ConstValue::Float(TWO_PI))?;
        let theta = u1.try_mul(&two_pi)?.cos()?;
        let one = u2.broadcast_scalar(ConstValue::Float(1.0))?;
        let neg_two = u2.broadcast_scalar(ConstValue::Float(-2.0))?;
        let r = one.try_sub(&u2)?.try_log()?.try_mul(&neg_two)?.try_sqrt()?;
        theta.try_mul(&r)
    }

    /// Normal `N(mean, std)` random tensor. Requires `std >= 0`.
    pub fn normal(shape: &[usize], mean: f64, std: f64) -> Result<Tensor> {
        if std < 0.0 {
            return Err(Error::ParamRange {
                op: "Tensor::normal",
                param: "std",
                value: format!("{std}"),
                constraint: ">= 0",
            });
        }
        let z = Tensor::randn(shape)?;
        let std_t = z.broadcast_scalar(ConstValue::Float(std))?;
        let mean_t = z.broadcast_scalar(ConstValue::Float(mean))?;
        z.try_mul(&std_t)?.try_add(&mean_t)
    }

    /// Uniform integer tensor `[low, high)`, dtype `int32`. Requires `low < high`.
    ///
    /// Truncates `(high - low) · rand` to int32 **before** adding `low`.
    /// Casting after the add would truncate-toward-zero asymmetrically for
    /// negative `low` (e.g. `low=-3, rand≈0.005` would yield `-2` instead of
    /// the correct `-3`).
    pub fn randint(shape: &[usize], low: i64, high: i64) -> Result<Tensor> {
        if low >= high {
            return Err(Error::ParamRange {
                op: "Tensor::randint",
                param: "low/high",
                value: format!("low={low}, high={high}"),
                constraint: "low < high",
            });
        }
        let scaled = Tensor::rand(shape)?;
        let range = scaled.broadcast_scalar(ConstValue::Float((high - low) as f64))?;
        let truncated = scaled.try_mul(&range)?.cast(DType::Int32)?;
        let offset = truncated.broadcast_scalar(ConstValue::Int(low))?;
        truncated.try_add(&offset)
    }

    /// `uniform(-1, 1) · prod(shape)^(-½)`. Same dtype contract as `uniform`.
    pub fn scaled_uniform(shape: &[usize]) -> Result<Tensor> {
        let numel: usize = shape.iter().copied().product::<usize>().max(1);
        let scale = (numel as f64).powf(-0.5);
        let u = Tensor::uniform(shape, -1.0, 1.0)?;
        let scale_t = u.broadcast_scalar(ConstValue::Float(scale))?;
        u.try_mul(&scale_t)
    }

    /// Glorot/Xavier uniform initializer, float32 output.
    pub fn glorot_uniform(shape: &[usize]) -> Result<Tensor> {
        Self::glorot_uniform_with_dtype(shape, DType::Float32)
    }

    /// Glorot/Xavier uniform initializer with explicit dtype.
    /// `bound = √(6 / (shape[0] + prod(shape[1..]))); uniform(-bound, bound)`.
    pub fn glorot_uniform_with_dtype(shape: &[usize], dtype: DType) -> Result<Tensor> {
        if shape.is_empty() {
            return Err(Error::ParamRange {
                op: "Tensor::glorot_uniform",
                param: "shape",
                value: "[]".to_string(),
                constraint: "at least 1D",
            });
        }
        let fan_in_v = fan_in(shape);
        let fan_out_v = shape[0];
        let bound = (6.0 / (fan_out_v + fan_in_v) as f64).sqrt();
        Self::uniform_with_dtype(shape, -bound, bound, dtype)
    }

    /// Kaiming/He uniform initializer for ReLU-family activations, float32 output.
    pub fn kaiming_uniform(shape: &[usize], a: f64) -> Result<Tensor> {
        Self::kaiming_uniform_with_dtype(shape, a, DType::Float32)
    }

    /// Kaiming/He uniform initializer with explicit dtype.
    ///
    /// `bound = √(6 / ((1 + a²) · prod(shape[1..]))); uniform(-bound, bound)`.
    ///
    /// `a` is the negative slope of the activation:
    /// - `0.0` — plain ReLU (PyTorch default).
    /// - `0.01` — leaky-ReLU with default slope.
    pub fn kaiming_uniform_with_dtype(shape: &[usize], a: f64, dtype: DType) -> Result<Tensor> {
        if shape.is_empty() {
            return Err(Error::ParamRange {
                op: "Tensor::kaiming_uniform",
                param: "shape",
                value: "[]".to_string(),
                constraint: "at least 1D",
            });
        }
        let bound = (6.0 / ((1.0 + a * a) * fan_in(shape) as f64)).sqrt();
        Self::uniform_with_dtype(shape, -bound, bound, dtype)
    }

    /// Kaiming/He normal initializer for ReLU-family activations.
    /// `std = √(2 / ((1 + a²) · prod(shape[1..]))); randn · std`.
    pub fn kaiming_normal(shape: &[usize], a: f64) -> Result<Tensor> {
        if shape.is_empty() {
            return Err(Error::ParamRange {
                op: "Tensor::kaiming_normal",
                param: "shape",
                value: "[]".to_string(),
                constraint: "at least 1D",
            });
        }
        let std = (2.0 / ((1.0 + a * a) * fan_in(shape) as f64)).sqrt();
        let z = Tensor::randn(shape)?;
        let std_t = z.broadcast_scalar(ConstValue::Float(std))?;
        z.try_mul(&std_t)
    }
}