Merge branch 'main' into jv-no-prefill

.github/workflows/test.yaml

@@ -5,110 +5,57 @@ on:
     branches:
       - main
 env:
-  OTP_VERSION: "25.0"
-  ELIXIR_VERSION: "1.14.0"
   MIX_ENV: test
   XLA_CACHE_DIR: ${{ github.workspace }}/cache/xla
   LIBTORCH_DIR: ${{ github.workspace }}/cache/torch
 jobs:
+  codespell:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: codespell-project/actions-codespell@v2
+        with:
+          skip: deps
+          ignore_words_list: whn,ehr
   main:
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - otp: "26.1.1"
+            elixir: "1.15.6"
+            lint: true
+          - otp: "25.3.2.6"
+            elixir: "1.14.5"
+          - otp: "25.3.2.6"
+            elixir: "1.14.5"
+            test_command_prepend: "USE_EXLA=true"
+          - otp: "25.3.2.6"
+            elixir: "1.14.5"
+            test_command_prepend: "USE_TORCHX=true"
     steps:
       - uses: actions/checkout@v3
       - name: Install Erlang & Elixir
         uses: erlef/setup-beam@v1
         with:
-          otp-version: "${{ env.OTP_VERSION }}"
-          elixir-version: "${{ env.ELIXIR_VERSION }}"
+          otp-version: "${{ matrix.otp }}"
+          elixir-version: "${{ matrix.elixir }}"
       - uses: actions/cache@v3
         with:
           path: |
             deps
             _build
             cache
-          key: ${{ runner.os }}-mix-${{ matrix.pair.elixir }}-${{ matrix.pair.otp }}-${{ hashFiles('**/mix.lock') }}
+          key: ${{ runner.os }}-mix-${{ matrix.elixir }}-${{ matrix.otp }}-${{ matrix.test_command_prepend }}-${{ hashFiles('**/mix.lock') }}
           restore-keys: |
             ${{ runner.os }}-mix-
       - name: Install mix dependencies
         run: mix deps.get
       - name: Check formatting
+        if: ${{ matrix.lint }}
         run: mix format --check-formatted
-      - name: Compile without optional deps
-        run: mix compile --skip-optional-deps --warnings-as-errors
+      - name: Check unused deps
+        if: ${{ matrix.lint }}
+        run: mix deps.unlock --check-unused
       - name: Run tests
-        run: mix test
-  exla_check:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v3
-      - name: Install Erlang & Elixir
-        uses: erlef/setup-beam@v1
-        with:
-          otp-version: "${{ env.OTP_VERSION }}"
-          elixir-version: "${{ env.ELIXIR_VERSION }}"
-      - uses: actions/cache@v3
-        with:
-          path: |
-            deps
-            _build
-            cache
-          key: ${{ runner.os }}-mix-${{ matrix.pair.elixir }}-${{ matrix.pair.otp }}-${{ hashFiles('**/mix.lock') }}
-          restore-keys: |
-            ${{ runner.os }}-mix-
-      - name: Install mix dependencies
-        run: mix deps.get
-      - name: Run tests against EXLA
-        run: USE_EXLA=true mix do compile --warnings-as-errors, test
-  torchx_check:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v3
-      - name: Install Erlang & Elixir
-        uses: erlef/setup-elixir@v1
-        with:
-          otp-version: "${{ env.OTP_VERSION }}"
-          elixir-version: "${{ env.ELIXIR_VERSION }}"
-      - uses: actions/cache@v3
-        with:
-          path: |
-            deps
-            _build
-            cache
-          key: ${{ runner.os }}-mix-${{ matrix.pair.elixir }}-${{ matrix.pair.otp }}-${{ hashFiles('**/mix.lock') }}
-          restore-keys: |
-            ${{ runner.os }}-mix-
-      - name: Install mix dependencies
-        run: mix deps.get
-      - name: Run tests against Torchx
-        run: USE_TORCHX=true mix do compile --warnings-as-errors, test
-  onnx_check:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Install Erlang & Elixir
-        uses: erlef/setup-beam@v1
-        with:
-          otp-version: "${{ env.OTP_VERSION }}"
-          elixir-version: "${{ env.ELIXIR_VERSION }}"
-      - name: Install Python
-        uses: actions/setup-python@v2
-        with:
-          python-version: "3.8"
-      - name: Install ONNX
-        run: pip install numpy onnx onnxruntime
-      - name: Install transformers
-        run: pip install git+https://github.com/huggingface/transformers.git sentencepiece pillow torch tensorflow
-      - name: Install Protoc
-        uses: arduino/setup-protoc@v1
-        with:
-          version: "3.x"
-      - name: Checkout AxonOnnx
-        uses: actions/checkout@v3
-        with:
-          repository: elixir-nx/axon_onnx
-          ref: refs/heads/master
-      - uses: actions/checkout@v3
-        with:
-          path: tmp/axon
-      - name: Run ONNX tests
-        run: |
-          AXON_PATH="tmp/axon" mix do deps.get, compile --warnings-as-errors, test
+        run: ${{ matrix.test_command_prepend }} mix do compile --skip-optional-deps --warnings-as-errors, test

.gitignore

@@ -25,6 +25,7 @@ axon-*.tar
 # Downloaded fixtures
 examples/vision/horses/
 examples/vision/humans/
+examples/structured/creditcard.csv
 
 # Temporary files for e.g. tests
 /tmp/

CHANGELOG.md

@@ -1,5 +1,52 @@
 # Changelog
 
+## v0.5.1 (2023-02-17)
+
+### Bug Fixes
+
+* Fixed incorrect results from group normalization
+
+## v0.5.0 (2023-02-16)
+
+### Enhancements
+
+* Bump Nx dependency
+* Update documentation to account for channels last default
+* Improve error message in compilation/build errors for models
+* Remove deprecated `transform`
+
+### Deprecations
+
+* Deprecate `Axon.Loop.handle/4`
+
+## v0.4.1 (2023-01-21)
+
+### Bug Fixes
+
+* Fixed a shape mismatch when training with certain optimizers
+
+## v0.4.0 (2023-01-19)
+
+### Enhancements
+
+* Add `Axon.pop_nodes/2` for popping nodes off of a graph
+* Update `Axon.freeze/2` and `Axon.unfreeze/2` for manipulating frozen portions of Axon graph
+* Add `Axon.Loop.monitor/5` for firing events based on loop state criteria
+* Add `Axon.Loop.kino_vega_lite_plot/4` for producing Kino plots during training
+* Add `Axon.Schedules.linear_decay/1`
+* Performance boosts to `Axon.Loop` which prevent compilation cache misses in most Axon training and evaluation loops
+* Add global event counts for more correct filtering during Axon loops
+* Use layer state to manage dropout keys, making training more deterministic when starting from the same key
+* Make building Axon models fully deterministic
+* Add a bidirectional combinator
+
+### Bug Fixes
+
+* Fix issue with namespaced stateful models not updating correctly during training
+* Fix bug in `Axon.Loop.early_stop/3` which incorrectly tracked progress and would not early stop loop
+* Fix bug in `Axon.Loop.reduce_lr_on_plateau/3` which incorrectly tracked progress and would not reduce learning rate
+* Fix bug in `Axon.Layers.conv_transpose/4` when using channels last
+
 ## v0.3.1 (2022-12-07)
 
 ### Enhancements
@@ -60,4 +107,4 @@
 
 ## v0.1.0 (2022-06-16)
 
-First release.
+First release.

README.md

@@ -8,10 +8,9 @@ Axon consists of the following components:
 
   * Functional API – A low-level API of numerical definitions (defn) of which all other APIs build on.
   * Model Creation API – A high-level model creation API which manages model initialization and application.
-  * Optimization API – An API for creating and using first-order optimization techniques based on the [Optax](https://github.com/deepmind/optax) library.
   * Training API – An API for quickly training models, inspired by [PyTorch Ignite](https://pytorch.org/ignite/index.html).
 
-Axon provides abstractions that enable easy integration while maintaining a level of separation between each component. You should be able to use any of the APIs without dependencies on others. By decoupling the APIs, Axon gives you full control over each aspect of creating and training a neural network.
+Axon provides abstractions that enable easy integration while maintaining a level of separation between each component. You should be able to use any of the APIs without dependencies on others. By decoupling the APIs, Axon gives you full control over each aspect of creating and training a neural network. Axon uses [Polaris](https://github.com/elixir-nx/polaris) for its optimization API.
 
 ## Overview
 
@@ -103,27 +102,13 @@ model =
 
 model_state =
   model
-  |> Axon.Loop.trainer(:categorical_cross_entropy, Axon.Optimizers.adamw(0.005))
+  |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(0.005))
   |> Axon.Loop.metric(:accuracy)
   |> Axon.Loop.handle(:iteration_completed, &log_metrics/1, every: 50)
   |> Axon.Loop.run(data, %{}, epochs: 10, compiler: EXLA)
 ```
 
-The step expects an optimizer as argument. The following are currently supported:
-
-* Adabelief
-* Adagrad
-* Adam
-* Adamw
-* Fromage
-* Lamb
-* Noisy SGD
-* Radam
-* RMSProp
-* SGD
-* Yogi
-
-It’s important to note that optimization API does not directly depend on Axon models. You can use the API to optimize any differentiable objective function.
+Axon uses [Polaris](https://github.com/elixir-nx/polaris) for its optimization API. It’s important to note that optimization API does not directly depend on Axon models. You can use the API to optimize any differentiable objective function.
 
 In the future we plan to support distributed training loops. We are also seeking ways to improve the performance of our training loops by running them entirely on native accelerators.
 
@@ -140,7 +125,7 @@ Then add Axon to your dependencies:
 ```elixir
 def deps do
   [
-    {:axon, "~> 0.2.0"}
+    {:axon, "~> 0.6"}
   ]
 end
 ```
@@ -150,13 +135,16 @@ You'll also likely want to include an `Nx` compiler such as `EXLA` for any pract
 ```elixir
 def deps do
   [
-    {:axon, "~> 0.2.0"},
-    {:exla, "~> 0.3.0"},
-    {:nx, "~> 0.3.0"}
+    {:axon, "~> 0.6"},
+    {:exla, "~> 0.6"},
   ]
 end
 ```
 
+## Integration with other platforms
+
+See [Ortex](https://github.com/elixir-nx/ortex) which provides full-blown compatibility for running ONNX models via ONNX Runtime bindings. Alternatively, see [AxonONNX](https://github.com/elixir-nx/axon_onnx) to convert ONNX models to Axon models whenever possible to achieve better integration with Nx.
+
 ## Sponsors
 
 <a href="https://dockyard.com"><img src="sponsors/dockyard.png" width=200 alt="DockYard"></a>

examples/basics/multi_input_example.exs

@@ -1,7 +1,7 @@
 Mix.install([
-  {:axon, "~> 0.3.0"},
-  {:exla, "~> 0.4.1"},
-  {:nx, "~> 0.4.1"}
+  {:axon, "~> 0.5"},
+  {:exla, "~> 0.5"},
+  {:nx, "~> 0.5"}
 ])
 
 defmodule XOR do
@@ -27,7 +27,7 @@ defmodule XOR do
   defp train_model(model, data, epochs) do
     model
     |> Axon.Loop.trainer(:binary_cross_entropy, :sgd)
-    |> Axon.Loop.run(data, %{}, epochs: epochs, iterations: 1000)
+    |> Axon.Loop.run(data, %{}, epochs: epochs, iterations: 1000, compiler: EXLA)
   end
 
   def run do

examples/basics/multi_output_example.exs

@@ -1,12 +1,9 @@
 Mix.install([
-  {:axon, "~> 0.3.0"},
-  # {:exla, "~> 0.2.2"},
-  {:nx, "~> 0.4.1"}
+  {:axon, "~> 0.5"},
+  {:exla, "~> 0.5"},
+  {:nx, "~> 0.5"}
 ])
 
-# Specify EXLA as the default defn compiler
-# EXLA.set_as_nx_default([:tpu, :cuda, :rocm, :host])
-
 defmodule Power do
   require Axon
 
@@ -34,11 +31,14 @@ defmodule Power do
     # model input and y is the target. Because we have multiple targets, we represent
     # y as a tuple. In the future, Axon will support any Nx container as an output
     data =
-      Stream.repeatedly(fn ->
-        # Batch size of 32
-        x = Nx.random_uniform({32, 1}, -10, 10, type: {:f, 32})
-        {x, {Nx.power(x, 2), Nx.power(x, 3)}}
-      end)
+      Stream.unfold(
+        Nx.Random.key(:erlang.system_time()),
+        fn key ->
+          # Batch size of 32
+          {x, next_key} = Nx.Random.uniform(key, -10, 10, shape: {32, 1}, type: {:f, 32})
+          {{x, {Nx.pow(x, 2), Nx.pow(x, 3)}}, next_key}
+        end
+      )
 
     # Create the training loop, notice we specify 2 MSE objectives, 1 for the first
     # output and 1 for the second output. This will create a loss function which is
@@ -62,7 +62,7 @@ defmodule Power do
     params =
       model
       |> Axon.Loop.trainer([mean_squared_error: 0.5, mean_squared_error: 0.5], :adam)
-      |> Axon.Loop.run(data, %{}, iterations: 250, epochs: 5)
+      |> Axon.Loop.run(data, %{}, iterations: 250, epochs: 5, compiler: EXLA)
 
     IO.inspect(Axon.predict(model, params, Nx.tensor([[3]])))
   end

examples/generative/fashionmnist_autoencoder.exs

@@ -1,16 +1,12 @@
 Mix.install([
-  {:axon, "~> 0.1.0"},
-  {:exla, "~> 0.2.2"},
-  {:nx, "~> 0.2.1"},
-  {:scidata, "~> 0.1.6"}
+  {:axon, "~> 0.5"},
+  {:exla, "~> 0.5"},
+  {:nx, "~> 0.5"},
+  {:scidata, "~> 0.1"}
 ])
 
-# Configure default platform with accelerator precedence as tpu > cuda > rocm > host
-EXLA.set_as_nx_default([:tpu, :cuda, :rocm, :host])
-
-defmodule Fashionmist do
+defmodule FashionMNIST do
   require Axon
-  alias Axon.Loop.State
 
   defmodule Autoencoder do
     defp encoder(x, latent_dim) do
@@ -22,7 +18,7 @@ defmodule Fashionmist do
     defp decoder(x) do
       x
       |> Axon.dense(784, activation: :sigmoid)
-      |> Axon.reshape({1, 28, 28})
+      |> Axon.reshape({:batch, 1, 28, 28})
     end
 
     def build_model(input_shape, latent_dim) do
@@ -37,7 +33,7 @@ defmodule Fashionmist do
     |> Nx.from_binary(type)
     |> Nx.reshape({elem(shape, 0), 1, 28, 28})
     |> Nx.divide(255.0)
-    |> Nx.to_batched_list(32)
+    |> Nx.to_batched(32)
   end
 
   defp train_model(model, train_images, epochs) do
@@ -58,7 +54,7 @@ defmodule Fashionmist do
 
     sample_image =
       train_images
-      |> hd()
+      |> Enum.fetch!(0)
       |> Nx.slice_along_axis(0, 1)
       |> Nx.reshape({1, 1, 28, 28})
 
@@ -71,4 +67,4 @@ defmodule Fashionmist do
   end
 end
 
-Fashionmist.run()
+FashionMNIST.run()