clausal-jax

JAX predicates for Clausal Prolog.

Wraps JAX as Clausal Prolog predicates: arrays (creation, math, shape operations, linear algebra, FFT), PRNG and sampling, function transforms, pytrees, sharding, jax.scipy, activations and initializers, and the optional Optax, Equinox and Flax layers.

Install

pip install clausal-jax

jax and jaxlib are pulled in as dependencies. Optax, Equinox and Flax are extras: pip install "clausal-jax[optax,equinox,flax]". Requires Python 3.13 or later.

Use

Import the predicates into a seam (.seam) module, as the package's own tests do. A Clausal Prolog (.clausal) module reaches Python only through the seam: put the imports in a .seam module and list it under [tool.clausal] python_bridges in your project's pyproject.toml (see Importing Prolog).

This test, from tests/fixtures/jax_array_tests.seam, shows the shape of a call:

-import_from(jax, [array, shape])

test("array from list") <- (
    array([1.0, 2.0, 3.0], A),
    shape(A, [3])
)

Documentation

• jax — JAX Array Operations
• jax_equinox — Equinox Layers, Filter Transforms, Pytree Utilities
• jax_flax — Flax Linen Models, Init, Apply
• jax_nn — JAX Activations and Initializers
• jax_optax — Optimisers, Schedules, Losses (Optax)
• jax_random — JAX PRNG and Sampling
• jax_scipy — JAX Special Functions and Distributions
• jax_sharding — JAX Device Placement and Sharding
• jax_transforms — JAX Function Transforms
• jax_tree — JAX Pytrees

License

MIT