Lux.jl

💻 Installation

import Pkg
Pkg.add("Lux")

[!TIP] To use Lux online, use Google Colab. The Julia Runtime comes pre-installed with Lux and Reactant!

Packages	Stable Version	Monthly Downloads	Total Downloads	Build Status
📦 Lux.jl
└ 📦 LuxLib.jl
└ 📦 LuxCore.jl

🚀 Benchmarks

Currently Benchmarks are scatter across a few places:

1. For comparison with other Julia packages like CUDA.jl take a look at Lux.jl/perf.
2. https://enzymead.github.io/Enzyme-JAX/benchmarks/ highlights performance of EnzymeJAX (backend for Reactant.jl) against JAX.
3. https://enzymead.github.io/Reactant.jl/benchmarks/ highlights performance of Reactant.jl against default XLA and base Julia compilation.

🤸 Quickstart

Reactant & Enzyme

using Lux, Random, Optimisers, Reactant, Enzyme

rng = Random.default_rng()
Random.seed!(rng, 0)

model = Chain(Dense(128, 256, tanh), Chain(Dense(256, 1, tanh), Dense(1, 10)))

dev = reactant_device()

ps, st = Lux.setup(rng, model) |> dev

x = rand(rng, Float32, 128, 2) |> dev

# We need to compile the model before we can use it.
model_forward = @compile model(x, ps, Lux.testmode(st))
model_forward(x, ps, Lux.testmode(st))

# Gradients can be computed using Enzyme
@jit Enzyme.gradient(Reverse, sum ∘ first ∘ Lux.apply, Const(model), x, ps, Const(st))

# All of this can be automated using the TrainState API
train_state = Training.TrainState(model, ps, st, Adam(0.001f0))

gs, loss, stats, train_state = Training.single_train_step!(
    AutoEnzyme(), MSELoss(),
    (x, dev(rand(rng, Float32, 10, 2))), train_state
)

Native Julia & Zygote

using Lux, Random, Optimisers, Zygote
# using LuxCUDA, AMDGPU, Metal, oneAPI # Optional packages for GPU support

# Seeding
rng = Random.default_rng()
Random.seed!(rng, 0)

# Construct the layer
model = Chain(Dense(128, 256, tanh), Chain(Dense(256, 1, tanh), Dense(1, 10)))

# Get the device determined by Lux
dev = gpu_device()

# Parameter and State Variables
ps, st = Lux.setup(rng, model) |> dev

# Dummy Input
x = rand(rng, Float32, 128, 2) |> dev

# Run the model
y, st = Lux.apply(model, x, ps, st)

# Gradients
## First construct a TrainState
train_state = Lux.Training.TrainState(model, ps, st, Adam(0.0001f0))

## We can compute the gradients using Training.compute_gradients
gs, loss, stats, train_state = Lux.Training.compute_gradients(AutoZygote(), MSELoss(),
    (x, dev(rand(rng, Float32, 10, 2))), train_state)

## Optimization
train_state = Training.apply_gradients!(train_state, gs) # or Training.apply_gradients (no `!` at the end)

# Both these steps can be combined into a single call
gs, loss, stats, train_state = Training.single_train_step!(AutoZygote(), MSELoss(),
    (x, dev(rand(rng, Float32, 10, 2))), train_state)

📚 Examples

Look in the examples directory for self-contained usage examples. The documentation has examples sorted into proper categories.

🆘 Getting Help

For usage related questions, please use Github Discussions which allows questions and answers to be indexed. To report bugs use github issues or even better send in a pull request.

🧑‍🔬 Citation

If you found this library to be useful in academic work, then please cite:

@software{pal2023lux,
  author    = {Pal, Avik},
  title     = {{Lux: Explicit Parameterization of Deep Neural Networks in Julia}},
  month     = apr,
  year      = 2023,
  note      = {If you use this software, please cite it as below.},
  publisher = {Zenodo},
  version   = {v1.4.2},
  doi       = {10.5281/zenodo.7808903},
  url       = {https://doi.org/10.5281/zenodo.7808903},
  swhid     = {swh:1:dir:1a304ec3243961314a1cc7c1481a31c4386c4a34;origin=https://doi.org/10.5281/zenodo.7808903;visit=swh:1:snp:e2bbe43b14bde47c4ddf7e637eb7fc7bd10db8c7;anchor=swh:1:rel:2c0c0ff927e7bfe8fc8bc43fd553ab392a6eb403;path=/}
}

@thesis{pal2023efficient,
  title     = {{On Efficient Training \& Inference of Neural Differential Equations}},
  author    = {Pal, Avik},
  year      = {2023},
  school    = {Massachusetts Institute of Technology}
}

Also consider starring our github repo.

🧑‍💻 Contributing

This section is somewhat incomplete. You can contribute by contributing to finishing this section 😜.

💎 Formatting (JuliaFormatter)

[!NOTE] Pin JuliaFormatter to v1 until upstream issues with v2 are resolved.

using JuliaFormatter
format(".")

🧪 Testing

The full test of Lux.jl takes a long time, here's how to test a portion of the code.

Tests are organized by directories, where each directory contains test files with @testset blocks. For example, tests for SkipConnection are in test/core_layers/containers_tests.jl.

Running a Specific Test File

The easiest way to run a specific test is to directly activate the test directory and include the test file:

# From the Lux.jl root directory
using Pkg
Pkg.activate("test")

# Run a specific test file
include("test/core_layers/containers_tests.jl")

This approach allows you to quickly iterate on specific tests without running the entire test suite.

See ParallelTestRunners.jl for details on executing specific groups of tests.

Running Test Groups via CI

To run a specific group of tests via the test runner, you can pass the directory name as a positional argument:

julia --project -e 'using Pkg; Pkg.test(test_args=["core_layers"])'

Running All Tests

To run the full test suite:

julia --project -e 'using Pkg; Pkg.test()'

📖 Documentation

Lux builds a bunch of tutorials as part of its documentation. This can be time-consuming and requires a lot of compute. To speed up the build, you can set the LUX_DOCS_DRAFT_BUILD=true.

LUX_DOCS_DRAFT_BUILD=true julia --threads=auto --startup=no --project=docs docs/make.jl

When writing tutorials (anything under examples/), include the tutorial in docs/tutorials.jl. If the tutorial is time-consuming, set should_run to false.

Additionally for a new page to be included in the navigation and sidebar, these need to be added to docs/src/.vitepress/config.mts. Specifically these need to be added under sidebar and/or nav based on the type of page.

To use LiveServer to preview the docs locally, checkout DocumenterVitepress.jl documentation.