LTNjax: JAX implementation of Logic Tensor Networks

Logic Tensor Networks (LTN)

LTNjax is a neurosymbolic framework that allows the implementation of knowledge in the form of logical expressions as objective for neural networks. LTN uses a differentiable first-order logic language, called Real Logic, to incorporate machine learning and logic.

This repository contains an implementation of LTN in Jax. This project is based on the projects logictensornetworks/logictensornetworks (TensorFlow 2) and tommasocarraro/LTNtorch (PyTorch).

Basics

The framework consists of the following components:

• Constants that are tensors of any rank $n \in [0,\infty)$, i.e. it is an element of $\mathbb{R}$ or an element of $\mathbb{R}^{(d_1 \times \dotsc \times d_n)}$ that we call feature dimensions. Constants can be trainable by putting nnx.Param values in.
• Variables that consists of $m$ individuals that are tensors of any rank. That means that an individual is an element of $\mathbb{R}^{(d_1 \times \dotsc \times d_n)}$ that we call feature dimensions. So a Variable is an element of $\mathbb{R}^{(m \times d_1 \times \dotsc \times d_n)}$. Variables can be trainable by putting nnx.Param values in.
• Predicates takes tuples of Constant and Variable objects as arguments and is applied to each configuration of the variables and returns a truth-value [0, 1]. Predicates consist of trainable nnx.Module objects that can be neural networks or lambda expressions.
• Functions takes tuples of Constant and Variable objects as arguments and is applied to each configuration of the variables and returns a tensor $\mathbb{R}^{(d_1 \times \dotsc \times d_n)}$. Functions consist of trainable nnx.Module objects that can be neural networks or lambda expressions.
• Connectives connects an arbitrary amount of Expressions to one joint Expressions, i.e. $\land, \lor, \neg, \Rightarrow, \leftrightarrow$.
• (Masked) Forall- or Exists-Quantifiers, i.e. $\forall, \exists$, that quantify over a set of Variables that satisfy a given mask and check whether they fulfill some given Expression.

Repository structure

• docs/ -- Documentation.
• examples/ -- Code examples.
• LICENSES/ -- All license files used somewhere in this project. Also see LICENSE.md.
• ltnjax/core.py -- Core system for defining constants, variables, functions, connectives and quantifiers.
• ltnjax/fuzzy_ops.py -- A collection of fuzzy logic operators defined using JAX primitives.
• scripts/ -- Scripts for testing the code.
• tests/ -- Tests.
• tutorials/ -- Tutorials written as jupyter notebooks.

Requirements

• Python 3.11 or later

• uv (highly recommended)

  # Linux and macOS
  curl -LsSf https://astral.sh/uv/install.sh | sh
  # Windows PowerShell
  powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex"

Install

There are multiple optional dependencies available:

• dev: installs additional development tools
• test: installs test requirements
• docs: installs documentation requirements
• all: installs all optional dependencies, except jax[cuda] and jax[tpu].
• jupyter-notebooks: installs requirements for the jupyter notebooks in tutorials/.

Of course, it is possible to install the software without any additional optional dependencies. Choose your poison based on your own requirements.

# update uv (optional)
uv self update

# create and install or update virtual environment with its dependencies
# choose one for the following depending on your system:
uv sync               # CPU
uv sync --group cuda  # GPU
uv sync --group tpu   # TPU

# activate the virtual environment (optional)
source .venv/bin/activate   # Linux and macOS
.venv\Scripts\activate.ps1  # Windows

# enable git pre-commit hooks (optional)
uv run prek install

Usage

import ltnjax as ltn

Getting Started

Tutorials

tutorials/ contains some important tutorials to getting started with coding in LTN. We suggest completing the tutorials in order. The tutorials cover the following topics:

1. Grounding in LTN (part 1): Real Logic, constants, predicates, functions, variables;
2. Grounding in LTN (part 2): connectives and quantifiers (+ complement: choosing appropriate operators for learning);
3. Learning in LTN: using satisfiability of LTN formulas as a training objective.

The tutorials are implemented using jupyter notebooks.

Training

To implement a trainable neural networks together with Logic Tensor Networks (LTN), the user can place a neural network as a Function/Predicate inside the LTN or put trainable values in Variables or Constants.

For that, define a nnx.Module and place modules (nnx.Module) and trainable parameters nnx.Param inside the __init__ function and the LTN inside the __call__ function.

For the first option, see train_mlp.py and for the second one see trainable_expr.py.

Examples

• The above example can be found in connectives.py.
• A basic example to train an MLP can be found in train_mlp.py.
• In the framework, every constant and variable can be trainable. For trainable values, just put nnx.Param as values and for non-trainable values, everything ArrayLike. See for example trainable_expr.py.
• During training, the user should take care about gradient issues, see gradient_issues.py.
• An example for masked quantifiers can be found in masked_quantifiers.py.
• In weighted_ops.py is illustrated how to use weighted aggregation operations.

Contribution

Please follow the contribution rules:

• use typed Python (type annotations)
• verify Python static code checks with ruff
• write tests
• document fixes, enhancements, new features, ...
• write scripts and examples OS independent or at least with linux, wsl support
• verify shell script static code check compliance with ShellCheck
• verify REUSE Specification 3.0 compliance for all files
• verify project license compliance without any license conflicts (e.g. for 3rd party libraries, data, models, ...)
• verify documentation (markdown) compliance w.r.t. markdown linting rules further specified inside the .markdownlint.json configuration file
• run all tests successfully

Documentation

This projects is using the Docstring style from Google. At least public classes, methods, fields, ... should be documented.

For further documentation we are using Markdown documentation with Material for MkDocs. See the docs folder for more details.

To locally serve the documentation, feel free to use:

uv run mkdocs serve

Testing

This projects provides a few different tests and checks.

Note

For detailed information about the tests, please check out the Testing section in the documentation.

In general, we have two files which defines all the tests:

• .pre-commit-config.yaml
• .gitlab-ci.yml

Install Node.js:

# Install Node.js
wget -qO- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.5/install.sh | bash
nvm install --lts

# Check version
nvm use --lts
# Check version
node --version
# v24.12.0

To easily run these tests locally (Node.js required), use:

./scripts/test.sh

If you only want to run the pre-commit hooks manually, use:

uv run prek run --all-files

Running the GitLab CI locally is a bit more complicated. Is also requires Node.js as well as Docker installed and configured.

npm exec gitlab-ci-local
# or run a single job, e.g. pre-commit
npm exec gitlab-ci-local -- pre-commit

Unittests

If you only want to run the unittests manually, use:

uv run unittest tests/test_core.py
uv run unittest tests/test_fuzzy_ops.py

Examples

To test the examples, use:

./scripts/test_examples.sh

Tutorials

To test the tutorials, use:

./scripts/test_tutorials.sh

Contributors

Citation

For accurate citation, refer to the corresponding metadata in the CITATION.cff file associated with this work.

License

Please see the file LICENSE for further information about how the content is licensed.

Funding

This open-source project was developed within the ROX project. This project has received public funding from the European Union NextGenerationEU within the Important Project of Common European Interest – Cloud Infrastructures and Services (IPCEI-CIS) under grant agreement 13IPC034.