jax-lnn 0.1.3

JAX Logical Neural Networks – neuro-symbolic framework with interval Łukasiewicz logic

JLNN – JAX Logical Neural Networks

---

Neuro-symbolic framework for interval-based Łukasiewicz logic built on JAX + Flax NNX.

JLNN enables turning symbolic logical rules into differentiable neural networks for training on data while maintaining interpretability and logical consistency.

Features

• Interval truth values [L, U]: Full support for uncertainty and contradiction modeling.
• Weighted logical gates: AND, OR, XOR, Implication, and NOT with Łukasiewicz semantics.
• Symbolic Compiler: Compile formulas like "0.8::A & B -> C" directly to NNX graphs.
• Temporal Logic: Experimental support for temporal operators (G, F, X).
• Logical Constraints: Built-in enforcement of axioms (e.g., weights $w \geq 1.0$).
• High Performance: JIT-compilation and hardware acceleration via JAX.
• Interoperability & Export: Native support for StableHLO and MLIR pipelines (recommended for production deployment). Export to ONNX and PyTorch is available with certain architectural limitations due to ONNX's lack of native structures for interval logic propagation.

Standard Installation (CPU Core)

# From PyPI
pip install jax-lnn[cpu,export]

# From GitHub (Bleeding-edge)
pip install git+https://github.com/RadimKozl/JLNN.git

Hardware Accelerated (Recommended)

# For NVIDIA GPU (CUDA 12+) and export pipelines (Colab T4/V100/A100)
pip install "jax-lnn[gpu,export]"

# For Google Cloud TPU support
pip install "jax-lnn[tpu,export]"

For Development

# For development
git clone https://github.com/RadimKozl/JLNN.git
cd JLNN

# Using uv (recommended for blistering speed)
uv sync --all-extras

# Using standard pip
pip install -e ".[test,docs,export,extra]"

⚠️ Google Colab Note: After running the installation on premium GPU instances (like A100), always restart your session gracefully via the UI menu (Runtime -> Restart session). Avoid using automated hard-kill lines like os.kill(os.getpid(), 9), as they crash virtualized XLA network sockets and cause browser WebSocket disconnections.

Quickstart

import jax
import jax.numpy as jnp
from flax import nnx
from jlnn.symbolic.compiler import LNNFormula
from jlnn.nn.constraints import apply_constraints
from jlnn.training.losses import total_lnn_loss, logical_mse_loss, contradiction_loss
from jlnn.storage.checkpoints import save_checkpoint, load_checkpoint
import optax

# 1. Define and compile the formula
model = LNNFormula("0.8::A & B -> C", nnx.Rngs(42))

# 2. Ground inputs (including initial state for C)
inputs = {
    "A": jnp.array([[0.9]]),
    "B": jnp.array([[0.7]]),
    "C": jnp.array([[0.5]])   # MANDATORY – consequent must have grounding!
}

target = jnp.array([[0.6, 0.85]])

# 3. Loss function
def loss_fn(model, inputs, target):
    pred = model(inputs)
    pred = jnp.nan_to_num(pred, nan=0.5, posinf=1.0, neginf=0.0)  # protection against NaN
    return total_lnn_loss(pred, target)

# 4. Initialize Optimizer
optimizer = nnx.Optimizer(
    model,
    wrt=nnx.Param,
    tx=optax.chain(
        optax.clip_by_global_norm(1.0),
        optax.adam(learning_rate=0.001)
    )
)

# 5. Training Step
@nnx.jit
def train_step(model, optimizer, inputs, target):
    # Gradients to the model – closure is traceable (inputs/target are arrays)
    grads = nnx.grad(lambda m: loss_fn(m, inputs, target))(model)

    # Loss before update (for debug)
    loss = loss_fn(model, inputs, target)

    optimizer.update(model, grads)
    apply_constraints(model)

    final_loss = loss_fn(model, inputs, target)
    final_pred = model(inputs)

    return loss, final_loss, final_pred

print("=== Starting training ===")
steps = 50
for step in range(steps):
    loss, final_loss, pred = train_step(model, optimizer, inputs, target)

    print(f"Step {step:3d} | Loss before/after constraints: {loss:.6f} → {final_loss:.6f}")
    print(f"Prediction: {pred}")
    print("─" * 60)

    if jnp.isnan(final_loss).any():
        print("❌ NaN detected! Stopping.")
        break

print("=== Training completed ===")

# 6. Result after training

final_pred = model(inputs)
print("\nFinal prediction after training:")
print(final_pred)

print(f"\nTarget interval: {target}")
print(f"Final loss: {total_lnn_loss(final_pred, target):.6f}")

See the introductory Jupyter notebook: Jax_lnn_base.ipynb

Acknowledgments & Inspiration

JLNN is inspired by and builds upon the foundations laid by several pioneering neuro-symbolic projects:

• LNN (IBM Research) – The primary inspiration for interval-based logical neural networks.
• LTNtorch – Logic Tensor Networks implementation in PyTorch.
• PyReason – Software for open-world temporal logic reasoning.

Documentation

• Online Documentation
• Introductory Tutorial
• API Reference
• GitHub
• PyPI

Discord channel

---

License

This project is licensed under the MIT License - see the LICENSE file for details.