JLNN Compiler

JLNN Symbolic Compiler

This module provides the infrastructure to compile symbolic logical formulas into neural computational graphs using Flax NNX and JAX. It leverages the Lark parser to transform strings into a tree of NNX Modules.

class jlnn.symbolic.compiler.BinaryGateNode(*args: Any, **kwargs: Any)

Bases: Node

Represents a gate with exactly 2 inputs, specifically designed for asymmetric operations like Implication (A -> B).

forward(values: Dict[str, Array]) → Array

Evaluates antecedent and consequent separately before applying the gate.

class jlnn.symbolic.compiler.JLNNCompiler(rngs: Rngs)

Bases: Transformer

Lark Transformer that converts a CST (Concrete Syntax Tree) into an NNX model tree.

Note

The grammar rules use and, or, and not. Since these are reserved keywords in Python, we define them as and_, or_, not_ and alias them dynamically after class definition.

and(children: List[Node]) → NAryGateNode

Constructs a WeightedAnd gate node.

and_(children: List[Node]) → NAryGateNode

Constructs a WeightedAnd gate node.

implication(children: List[Node]) → BinaryGateNode

Constructs a WeightedImplication gate node (A -> B).

not(children: List[Node]) → UnaryGateNode

Constructs a WeightedNot gate node.

not_(children: List[Node]) → UnaryGateNode

Constructs a WeightedNot gate node.

or(children: List[Node]) → NAryGateNode

Constructs a WeightedOr gate node.

or_(children: List[Node]) → NAryGateNode

Constructs a WeightedOr gate node.

variable(tokens: List[Token]) → PredicateNode

Transforms a variable token into a PredicateNode, ensuring weight sharing if name repeats.

weighted_expr(children: List[Any]) → Node

Root rule processor; returns the final compiled node structure.

class jlnn.symbolic.compiler.LNNFormula(*args: Any, **kwargs: Any)

Bases: Module

The main LNN Formula wrapper. Compiles a logical string formula into a differentiable NNX module.

root

The root node of the neural computational graph.

predicates

An nnx.Dict containing all trainable predicate groundings.

Example

>>> rngs = nnx.Rngs(0)
>>> model = LNNFormula("A & B -> C", rngs)
>>> inputs = {"A": jnp.array([[0.9]]), "B": jnp.array([[0.8]]), "C": jnp.array([[0.1]])}
>>> output = model(inputs)  # Returns truth interval [Lower, Upper]

class jlnn.symbolic.compiler.NAryGateNode(*args: Any, **kwargs: Any)

Bases: Node

Represents a logic gate with N inputs, such as weighted AND, OR or XOR.

forward(values: Dict[str, Array]) → Array

Evaluates all children, stacks results, and applies the gate logic.

class jlnn.symbolic.compiler.Node(*args: Any, **kwargs: Any)

Bases: Module

Abstract base class for all nodes in the JLNN computational graph.

Each node represents a logical operation or a predicate and must implement a forward pass that operates on JAX arrays.

forward(values: Dict[str, Array]) → Array

Performs the forward evaluation of the node.

Parameters:

values – A mapping of variable names to JAX arrays containing the raw input features.

Returns:

A JAX array representing the truth value (typically a truth interval).

class jlnn.symbolic.compiler.PredicateNode(*args: Any, **kwargs: Any)

Bases: Node

Represents a leaf node (variable) mapping to a LearnedPredicate.

This node acts as the grounding layer where raw numeric data is transformed into a fuzzy truth value.

forward(values: Dict[str, Array]) → Array

Retrieves input data by name and passes it through the predicate.

class jlnn.symbolic.compiler.UnaryGateNode(*args: Any, **kwargs: Any)

Bases: Node

Represents a gate with a single input, such as NOT.

forward(values: Dict[str, Array]) → Array

Passes the child’s output through the unary gate.

The compiler recursively traverses the syntax tree from the parser and creates corresponding NNX modules for each node. The result is a hierarchical model that can be trained directly.

Key Components:

• Node: Abstract base for all graph nodes. Ensures recursive forward calls.

• PredicateNode: Represents input variables (leaves of the graph).

• NAryGateNode: Zapouzdřuje n-ární hradla (AND, OR, XOR).

Main class that unifies the process from a string formula to an executable model.