tensorlogic-compiler 0.1.0-rc.1

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
139
140
141
142
143
144
145
146
147
148
149
150

//! Compilation of implication (→).

use anyhow::{bail, Result};
use tensorlogic_ir::{EinsumGraph, EinsumNode, TLExpr};

use crate::context::{CompileState, CompilerContext};

use super::compile_expr;

pub(crate) fn compile_imply(
    premise: &TLExpr,
    conclusion: &TLExpr,
    ctx: &mut CompilerContext,
    graph: &mut EinsumGraph,
) -> Result<CompileState> {
    let mut premise_state = compile_expr(premise, ctx, graph)?;
    let mut conclusion_state = compile_expr(conclusion, ctx, graph)?;

    // Align axes between premise and conclusion
    // Strategy:
    // 1. If premise has extra axes → marginalize them out (∀ quantification)
    // 2. If conclusion has extra axes → broadcast premise to match (independent axes)
    // 3. Result has union of all axes from both operands

    if premise_state.axes != conclusion_state.axes {
        // Find axes in premise not in conclusion (need to marginalize)
        let premise_extra: Vec<char> = premise_state
            .axes
            .chars()
            .filter(|c| !conclusion_state.axes.contains(*c))
            .collect();

        // Find axes in conclusion not in premise (need to broadcast)
        let conclusion_extra: Vec<char> = conclusion_state
            .axes
            .chars()
            .filter(|c| !premise_state.axes.contains(*c))
            .collect();

        // Marginalize extra premise axes
        if !premise_extra.is_empty() {
            let axes_to_reduce: Vec<usize> = premise_extra
                .iter()
                .map(|c| premise_state.axes.find(*c).unwrap())
                .collect();

            let reduce_name = ctx.fresh_temp();
            let reduce_idx = graph.add_tensor(reduce_name);
            let reduce_node =
                EinsumNode::reduce("sum", axes_to_reduce, premise_state.tensor_idx, reduce_idx);
            graph.add_node(reduce_node)?;

            let new_axes: String = premise_state
                .axes
                .chars()
                .filter(|c| !premise_extra.contains(c))
                .collect();

            premise_state = CompileState {
                tensor_idx: reduce_idx,
                axes: new_axes,
            };
        }

        // Broadcast premise to include conclusion's extra axes
        if !conclusion_extra.is_empty() {
            // Use einsum to broadcast: "ab->abc" where c is the extra axis
            let mut output_axes = premise_state.axes.clone();
            for &c in &conclusion_extra {
                output_axes.push(c);
            }

            let broadcast_spec = format!("{}->{}", premise_state.axes, output_axes);
            let broadcast_name = ctx.fresh_temp();
            let broadcast_idx = graph.add_tensor(broadcast_name);
            let broadcast_node = EinsumNode::new(
                broadcast_spec,
                vec![premise_state.tensor_idx],
                vec![broadcast_idx],
            );
            graph.add_node(broadcast_node)?;

            premise_state = CompileState {
                tensor_idx: broadcast_idx,
                axes: output_axes,
            };
        }

        // Similarly broadcast conclusion if needed
        let premise_remaining: Vec<char> = premise_state
            .axes
            .chars()
            .filter(|c| !conclusion_state.axes.contains(*c))
            .collect();

        if !premise_remaining.is_empty() {
            let mut output_axes = conclusion_state.axes.clone();
            for &c in &premise_remaining {
                output_axes.push(c);
            }

            let broadcast_spec = format!("{}->{}", conclusion_state.axes, output_axes);
            let broadcast_name = ctx.fresh_temp();
            let broadcast_idx = graph.add_tensor(broadcast_name);
            let broadcast_node = EinsumNode::new(
                broadcast_spec,
                vec![conclusion_state.tensor_idx],
                vec![broadcast_idx],
            );
            graph.add_node(broadcast_node)?;

            conclusion_state = CompileState {
                tensor_idx: broadcast_idx,
                axes: output_axes,
            };
        }

        // Final check: axes should now be aligned
        if premise_state.axes != conclusion_state.axes {
            bail!(
                "IMPLY operands have incompatible axes after alignment: '{}' vs '{}'",
                premise_state.axes,
                conclusion_state.axes
            );
        }
    }

    // Implication: a → b == ReLU(b - a)
    // Step 1: Subtract (conclusion - premise)
    let subtract_name = ctx.fresh_temp();
    let subtract_idx = graph.add_tensor(subtract_name);
    let subtract_node = EinsumNode::elem_binary(
        "subtract",
        conclusion_state.tensor_idx,
        premise_state.tensor_idx,
        subtract_idx,
    );
    graph.add_node(subtract_node)?;

    // Step 2: Apply ReLU
    let result_name = ctx.fresh_temp();
    let result_idx = graph.add_tensor(result_name);
    let relu_node = EinsumNode::elem_unary("relu", subtract_idx, result_idx);
    graph.add_node(relu_node)?;

    Ok(CompileState {
        tensor_idx: result_idx,
        axes: premise_state.axes,
    })
}