tensorlogic-compiler 0.1.0

Compiler for transforming logic expressions into tensor 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
137
138

//! Helper constructors and the algebraic simplification pass.

use tensorlogic_ir::TLExpr;

/// Build the canonical zero expression.
#[inline]
pub(super) fn zero() -> TLExpr {
    TLExpr::Constant(0.0)
}

/// Build the canonical one expression.
#[inline]
pub(super) fn one() -> TLExpr {
    TLExpr::Constant(1.0)
}

/// Check whether `expr` is a constant equal to `value` (within f64 epsilon).
#[inline]
pub(super) fn is_constant_value(expr: &TLExpr, value: f64) -> bool {
    match expr {
        TLExpr::Constant(v) => (v - value).abs() < f64::EPSILON,
        _ => false,
    }
}

/// Return a symbolic representation of the derivative of a function node.
///
/// For a known named function `f` (zero-arity predicate), the marker is `Pred("d_f", [])`.
/// For a complex/anonymous function, `Pred("d_f", [])` is used as a generic marker.
pub(super) fn derivative_of_function(function: &TLExpr) -> TLExpr {
    match function {
        TLExpr::Pred { name, args } if args.is_empty() => {
            TLExpr::pred(format!("d_{}", name), vec![])
        }
        _ => TLExpr::pred("d_f".to_string(), vec![]),
    }
}

/// Basic algebraic simplification applied to derivative expressions.
///
/// Simplification rules:
/// - `0 + x → x`,  `x + 0 → x`
/// - `0 * x → 0`,  `x * 0 → 0`
/// - `1 * x → x`,  `x * 1 → x`
/// - `x - 0 → x`
/// - `0 - c → -c` (constant folding for arithmetic negation form)
/// - `0 / x → 0`
/// - `x ^ 0 → 1`,  `x ^ 1 → x`
/// - constant folding for pure-constant arithmetic nodes
pub fn simplify_derivative(expr: TLExpr) -> TLExpr {
    match expr {
        TLExpr::Add(l, r) => {
            let l = simplify_derivative(*l);
            let r = simplify_derivative(*r);
            if is_constant_value(&l, 0.0) {
                return r;
            }
            if is_constant_value(&r, 0.0) {
                return l;
            }
            if let (TLExpr::Constant(a), TLExpr::Constant(b)) = (&l, &r) {
                return TLExpr::Constant(a + b);
            }
            TLExpr::Add(Box::new(l), Box::new(r))
        }

        TLExpr::Sub(l, r) => {
            let l = simplify_derivative(*l);
            let r = simplify_derivative(*r);
            if is_constant_value(&r, 0.0) {
                return l;
            }
            if let (TLExpr::Constant(a), TLExpr::Constant(b)) = (&l, &r) {
                return TLExpr::Constant(a - b);
            }
            TLExpr::Sub(Box::new(l), Box::new(r))
        }

        TLExpr::Mul(l, r) => {
            let l = simplify_derivative(*l);
            let r = simplify_derivative(*r);
            if is_constant_value(&l, 0.0) || is_constant_value(&r, 0.0) {
                return TLExpr::Constant(0.0);
            }
            if is_constant_value(&l, 1.0) {
                return r;
            }
            if is_constant_value(&r, 1.0) {
                return l;
            }
            if let (TLExpr::Constant(a), TLExpr::Constant(b)) = (&l, &r) {
                return TLExpr::Constant(a * b);
            }
            TLExpr::Mul(Box::new(l), Box::new(r))
        }

        TLExpr::Div(l, r) => {
            let l = simplify_derivative(*l);
            let r = simplify_derivative(*r);
            if is_constant_value(&l, 0.0) {
                return TLExpr::Constant(0.0);
            }
            if let (TLExpr::Constant(a), TLExpr::Constant(b)) = (&l, &r) {
                if b.abs() > f64::EPSILON {
                    return TLExpr::Constant(a / b);
                }
            }
            TLExpr::Div(Box::new(l), Box::new(r))
        }

        TLExpr::Pow(base, exp) => {
            let base = simplify_derivative(*base);
            let exp = simplify_derivative(*exp);
            if is_constant_value(&exp, 0.0) {
                return TLExpr::Constant(1.0);
            }
            if is_constant_value(&exp, 1.0) {
                return base;
            }
            if let (TLExpr::Constant(b), TLExpr::Constant(e)) = (&base, &exp) {
                return TLExpr::Constant(b.powf(*e));
            }
            TLExpr::Pow(Box::new(base), Box::new(exp))
        }

        TLExpr::And(l, r) => TLExpr::And(
            Box::new(simplify_derivative(*l)),
            Box::new(simplify_derivative(*r)),
        ),
        TLExpr::Or(l, r) => TLExpr::Or(
            Box::new(simplify_derivative(*l)),
            Box::new(simplify_derivative(*r)),
        ),
        TLExpr::Not(inner) => TLExpr::Not(Box::new(simplify_derivative(*inner))),

        other => other,
    }
}