minidx-vis 0.1.10

Visualization crate for minidx, a simple + compile-time neural-network 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
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

use crate::{PaintParams, ParamVisOpts};
use minidx_core::{Dtype, Float};
use raqote::DrawTarget;

/// Identifies layer types which are native to minidx: needed to get around
/// trait conflicts between impls on M and (M,)
trait LayerMarker {}

impl<E: Dtype, const I: usize> LayerMarker for minidx_core::layers::Bias1d<E, I> {}
impl<E: Dtype + minidx_core::matmul::MatMulImpl, const I: usize, const O: usize> LayerMarker
    for minidx_core::layers::Dense<E, I, O>
{
}
impl<E: Float> LayerMarker for minidx_core::layers::Activation<E> {}
impl<E: Dtype> LayerMarker for minidx_core::layers::ScalarScale<E> {}
impl<E: Float, const I: usize> LayerMarker for minidx_core::layers::Swish<E, I> {}
impl LayerMarker for minidx_core::layers::Softmax {}
impl<E: Dtype, const I: usize> LayerMarker for minidx_core::layers::Diag<E, I> {}
impl<E: Float + minidx_core::matmul::MatMulImpl, const I: usize> LayerMarker
    for minidx_core::layers::RMSDiv<E, I>
{
}

impl<
        E: Dtype,
        const I: usize,
        M: Default + minidx_core::Module<[E; I]> + minidx_core::VisualizableUnit,
    > LayerMarker for minidx_core::layers::LR<E, I, M>
{
}

/// Some composition of parameters which can be visualized.
pub trait VisualizableNetwork<DT> {
    type Params: std::fmt::Debug + Sized;

    fn visualize(&self, dt: &mut DT, opts: &mut ParamVisOpts) -> (f32, f32);
}

impl<M: minidx_core::VisualizableUnit + LayerMarker> VisualizableNetwork<DrawTarget> for M
where
    DrawTarget: PaintParams<M::Params>,
{
    type Params = M::Params;

    fn visualize(&self, dt: &mut DrawTarget, opts: &mut ParamVisOpts) -> (f32, f32) {
        let bounds = dt.layout_bounds(&opts);
        dt.paint_params(self.params(), opts);
        (0.0, bounds.1)
    }
}

impl<
        E: Dtype + Float + minidx_core::matmul::MatMulImpl,
        const I: usize,
        const O: usize,
        A: minidx_core::Module<[E; O], Output = [E; O]> + Default,
    > VisualizableNetwork<DrawTarget> for minidx_core::layers::GLU<E, I, O, A>
where
    DrawTarget: PaintParams<[[E; O]; 1]> + PaintParams<[[E; I]; O]>,
{
    type Params = (); // Not technically correct but it'll do

    fn visualize(&self, dt: &mut DrawTarget, opts: &mut ParamVisOpts) -> (f32, f32) {
        let (gc, gb, sc, sb) = self.connection_params();

        let gc_box = <raqote::DrawTarget as PaintParams<[[E; I]; O]>>::layout_bounds(dt, &opts);
        dt.paint_params(gc, opts);
        let gb_box = <raqote::DrawTarget as PaintParams<[[E; O]; 1]>>::layout_bounds(
            dt,
            opts.update_cursor((0.0, gc_box.1)),
        );
        dt.paint_params(
            unsafe {
                // SAFETY: An array of N is exactly the same as a unary array of the array of N
                std::mem::transmute::<&[E; O], &[[E; O]; 1]>(gb)
            },
            opts,
        );

        let sc_box = <raqote::DrawTarget as PaintParams<[[E; I]; O]>>::layout_bounds(
            dt,
            opts.update_cursor((0.0, gb_box.1)),
        );
        dt.paint_params(sc, opts);
        let sb_box = <raqote::DrawTarget as PaintParams<[[E; O]; 1]>>::layout_bounds(
            dt,
            opts.update_cursor((0.0, sc_box.1)),
        );
        dt.paint_params(
            unsafe {
                // SAFETY: An array of N is exactly the same as a unary array of the array of N
                std::mem::transmute::<&[E; O], &[[E; O]; 1]>(sb)
            },
            opts,
        );

        (0.0, sb_box.1)
    }
}

macro_rules! tuple_impls {
    ([$($name:ident),+] [$($idx:tt),*], $last:ident, [$($rev_tail:ident),*]) => {
        impl<
            $last:
            $(VisualizableNetwork<DrawTarget>, $rev_tail: )*
            VisualizableNetwork<DrawTarget>
        > VisualizableNetwork<DrawTarget> for ($($name,)+) {
            type Params = (
                $($name::Params,)+
            );

            fn visualize(&self, dt: &mut DrawTarget, opts: &mut ParamVisOpts) -> (f32, f32) {
                let bounds = self.0.visualize(dt, opts);
                $(let bounds = self.$idx.visualize(dt, opts.update_cursor(bounds));)*
                bounds
            }
        }
    }
}

tuple_impls!([M1][], M1, []);
tuple_impls!([M1, M2][1], M2, [M1]);
tuple_impls!([M1, M2, M3] [1, 2], M3, [M2, M1]);
tuple_impls!([M1, M2, M3, M4] [1, 2, 3], M4, [M3, M2, M1]);
tuple_impls!([M1, M2, M3, M4, M5] [1, 2, 3, 4], M5, [M4, M3, M2, M1]);
tuple_impls!([M1, M2, M3, M4, M5, M6] [1, 2, 3, 4, 5], M6, [M5, M4, M3, M2, M1]);
tuple_impls!([M1, M2, M3, M4, M5, M6, M7] [1, 2, 3, 4, 5, 6], M7, [M6, M5, M4, M3, M2, M1]);