twisterl 0.5.1

Reinforcement learning primitives and a Python extension for high performance training and inference.
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

// -*- coding: utf-8 -*-
/* 
(C) Copyright 2025 IBM. All Rights Reserved.

This code is licensed under the Apache License, Version 2.0. You may
obtain a copy of this license in the LICENSE.txt file in the root directory
of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.

Any modifications or derivative works of this code must retain this
copyright notice, and modified files need to carry a notice indicating
that they have been altered from the originals.
*/

use nalgebra::{DMatrix, DVector};

#[derive(Clone)]
pub struct Linear {
    weights: DMatrix<f32>,
    bias: DVector<f32>,
    apply_relu: bool
}


impl Linear {
    pub fn new(weights_vector: Vec<f32>, bias_vector: Vec<f32>, apply_relu: bool) -> Self {
        let weights = DMatrix::from_vec(bias_vector.len(), weights_vector.len() / bias_vector.len(), weights_vector);
        let bias = DVector::from_vec(bias_vector);
        Self { weights, bias, apply_relu }
    }
    
    pub fn forward(&self, input: &DVector<f32>) -> DVector<f32> {
        let mut out = (&self.weights * input) + &self.bias;
        if self.apply_relu {
            out =  out.map(relu);
        }
        out
    }
}

#[derive(Clone)]
pub struct EmbeddingBag {
    vectors: Vec<DVector<f32>>,
    bias: DVector<f32>,
    apply_relu: bool,
    obs_shape: Vec<usize>,
    conv_dim: usize
}

impl EmbeddingBag {
    pub fn new(vec_vectors: Vec<Vec<f32>>, bias_vector: Vec<f32>, apply_relu: bool, obs_shape: Vec<usize>, conv_dim: usize) -> Self {
        let vectors = vec_vectors.into_iter().map(|vec| DVector::from_vec(vec)).collect();
        let bias = DVector::from_vec(bias_vector);
        Self { vectors, bias, apply_relu, obs_shape, conv_dim }
    }

    pub fn forward(&self, input: &Vec<usize>) -> DVector<f32> {
        let mut out = self.bias.clone();
        if self.obs_shape.len() == 1 {
            // This is standard embeddings / linear
            for &i in input.iter() {
                out += &self.vectors[i];
            }
        } else if self.obs_shape.len() == 2 {
            let v_size = self.vectors[0].len();
            // This is conv1d
            for &i in input.iter() {
                // obs_shape[0] is the size of each col (i.e. the number of rows)
                let mut row = i / self.obs_shape[1]; 
                let mut col = i % self.obs_shape[1];

                // If conv_dim is 1 then we swap row and col
                if self.conv_dim == 1 {(row, col) = (col, row);}

                let mut out_slice = out.rows_mut(col * v_size, v_size);
                out_slice += &self.vectors[row];

            }

        } // TODO: add conv2d


        if self.apply_relu {
            out =  out.map(relu);
        }
        out
    }
}

fn relu(x: f32) -> f32 {
   if x > 0.0 { x } else { 0.0 }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_linear_forward() {
        let linear = Linear::new(vec![1.0, 2.0, 3.0, 4.0], vec![1.0, 1.0], false);
        let input = DVector::from_vec(vec![1.0, 2.0]);
        let out = linear.forward(&input);
        assert_eq!(out, DVector::from_vec(vec![8.0, 11.0]));
    }

    #[test]
    fn test_linear_forward_relu() {
        let linear = Linear::new(vec![-1.0, -2.0, 0.0, 1.0], vec![0.0, 0.0], true);
        let input = DVector::from_vec(vec![1.0, 2.0]);
        let out = linear.forward(&input);
        assert_eq!(out, DVector::from_vec(vec![0.0, 0.0]));
    }

    #[test]
    fn test_embedding_bag_forward() {
        let emb = EmbeddingBag::new(
            vec![vec![1.0, 2.0], vec![3.0, 4.0]],
            vec![0.0, 0.0],
            false,
            vec![2],
            0,
        );
        let out = emb.forward(&vec![0, 1]);
        assert_eq!(out, DVector::from_vec(vec![4.0, 6.0]));
    }
}