Skip to main content

matmul

Function matmul 

Source
pub fn matmul(builder: &Builder, f: Var, g: Var) -> Var
Expand description

Batch matmul

Examples found in repository?
examples/hidden.rs (line 122)
110    fn def(&self, builder: &Builder, [x]: [Var; 1]) -> [Var; 1] {
111        // Flatten input from B×28×28 to B×784
112        let [batch_size, h, w] = unpack::<3>(builder, shape(builder, x.clone()));
113        let flat_size = h * w;
114        let flat_shape = pack::<2>(builder, [batch_size, flat_size]);
115        let x = reshape(builder, flat_shape, x);
116
117        let root = self.path();
118
119        let p = param(builder, &root.extend(["0", "weights"]).unwrap());
120
121        // layer 1: B×784 @ 784×100 = B×100
122        let x = matmul(builder, x, p);
123        let x = nn::Sigmoid.call(builder, [x]);
124
125        // layer 2: B×100 @ 100×10 = B×10
126        let p = param(builder, &root.extend(["1", "weights"]).unwrap());
127        let x = matmul(builder, x, p);
128        let x = nn::Sigmoid.call(builder, [x]);
129
130        // result
131        [x]
132    }