tidu 0.2.0

Automatic-differentiation transforms (linearize, transpose, eager reverse-mode) for primitive computation graphs.
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

use computegraph::{GraphOperation, LocalValueId, OperationRole, ValueKey, ValueRef};
use std::hint::black_box;
use tidu::rules::{
    ADKey as ModuleADKey, ADRuleError as ModuleADRuleError, ADRuleKind as ModuleADRuleKind,
    ADRuleResult as ModuleADRuleResult, DiffPassId as ModuleDiffPassId,
    Primitive as ModulePrimitive,
};
use tidu::{
    ADKey, ADRuleError, ADRuleKind, ADRuleResult, DiffPassId, Primitive, PrimitiveBuilder,
    PrimitiveValue,
};

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
enum Key {
    Base(&'static str),
    Tangent { of: Box<Key>, pass: DiffPassId },
}

impl ADKey for Key {
    fn tangent_of(&self, pass: DiffPassId) -> Self {
        Self::Tangent {
            of: Box::new(self.clone()),
            pass,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct AddOp;

impl GraphOperation for AddOp {
    type Operand = f64;
    type Context = ();
    type InputKey = Key;

    fn input_count(&self) -> usize {
        2
    }

    fn output_count(&self) -> usize {
        1
    }
}

impl Primitive for AddOp {
    type ADContext = ();

    fn add() -> Self {
        Self
    }

    fn jvp_rule(
        &self,
        _builder: &mut impl PrimitiveBuilder<Self>,
        _primal_inputs: &[ValueKey<Self>],
        _primal_outputs: &[ValueKey<Self>],
        tangent_inputs: &[Option<LocalValueId>],
        _ctx: &mut Self::ADContext,
    ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
        Ok(vec![tangent_inputs[0].or(tangent_inputs[1])])
    }

    fn transpose_rule(
        &self,
        _builder: &mut impl PrimitiveBuilder<Self>,
        cotangent_outputs: &[Option<LocalValueId>],
        _inputs: &[PrimitiveValue<Self>],
        _mode: &OperationRole,
        _ctx: &mut Self::ADContext,
    ) -> tidu::ADRuleResult<Vec<Option<LocalValueId>>> {
        Ok(vec![cotangent_outputs[0], cotangent_outputs[0]])
    }
}

#[test]
fn root_reexports_match_rules_module_contract() {
    fn assert_key<K: ADKey + ModuleADKey>() {}
    fn assert_primitive<Op: Primitive + ModulePrimitive>()
    where
        Op::InputKey: ADKey,
    {
    }
    fn assert_result<T>(result: ADRuleResult<T>) -> ModuleADRuleResult<T> {
        result
    }
    fn assert_pass_id(pass: DiffPassId) -> ModuleDiffPassId {
        pass
    }

    assert_key::<Key>();
    assert_primitive::<AddOp>();
    let tangent = Key::Base("x").tangent_of(7);
    assert!(matches!(tangent, Key::Tangent { .. }));
    assert_eq!(assert_pass_id(7), 7);
    assert_eq!(ModuleADRuleKind::Jvp.as_str(), ADRuleKind::Jvp.as_str());
    assert_eq!(ModuleADRuleKind::Transpose.as_str(), "transpose");

    let err: ADRuleError = ModuleADRuleError::unsupported("test::op", ModuleADRuleKind::Jvp);
    assert_eq!(err.to_string(), "unsupported jvp AD rule for test::op");
    assert!(std::error::Error::source(&err).is_none());
    let rule_fn: fn(&ADRuleError) -> ADRuleKind = ADRuleError::rule;
    let runtime_err = black_box(assert_result::<()>(Err(err)).unwrap_err());
    assert_eq!(rule_fn(&runtime_err), ADRuleKind::Jvp);

    let transpose_err =
        ModuleADRuleError::unsupported("test::transpose", ModuleADRuleKind::Transpose);
    assert_eq!(
        transpose_err.to_string(),
        "unsupported transpose AD rule for test::transpose"
    );

    let invalid: ADRuleError = ModuleADRuleError::invalid_input(
        "test::solve",
        ModuleADRuleKind::Transpose,
        "expected rank >= 2",
    );
    assert_eq!(rule_fn(&invalid), ADRuleKind::Transpose);
    assert_eq!(
        invalid.to_string(),
        "invalid transpose AD input for test::solve: expected rank >= 2"
    );
}

#[test]
fn primitive_value_round_trips_computegraph_value_refs() {
    let local = PrimitiveValue::<AddOp>::Local(3);
    let local_ref: ValueRef<AddOp> = local.clone().into();
    assert_eq!(local_ref, ValueRef::Local(3));
    assert_eq!(PrimitiveValue::from(local_ref), local);

    let key = ValueKey::Input(Key::Base("x"));
    let external = PrimitiveValue::<AddOp>::External(key.clone());
    let external_ref: ValueRef<AddOp> = external.clone().into();
    assert_eq!(external_ref, ValueRef::External(key));
    assert_eq!(PrimitiveValue::from(external_ref), external);
}