chore: Update to Nx, EXLA and Torchx 1.0

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## Unreleased
+
+### Enhancements
+
+* Depend on the released Nx, EXLA and Torchx 1.0 instead of Nx's main branch
+
 ## v0.8.1 (2026-03-11)
 
 ### Bug Fixes

lib/axon.ex

@@ -3978,7 +3978,7 @@ defmodule Axon do
   @doc """
   Attaches a hook to the given Axon model.
 
-  Hooks compile down to `Nx.Defn.Kernel.hook/3` and provide the same
+  Hooks compile down to `Nx.Defn.Kernel.io_call/3` and provide the same
   functionality for adding side-effecting operations to a compiled
   model. For example, you can use hooks to inspect intermediate activations,
   send data to an external service, and more.

mix.exs

@@ -41,9 +41,9 @@ defmodule Axon.MixProject do
   # Run "mix help deps" to learn about dependencies.
   defp deps do
     [
-      {:nx, "~> 0.13.1", nx_opts()},
-      {:exla, "~> 0.13", [only: :test] ++ exla_opts()},
-      {:torchx, "~> 0.13", [only: :test] ++ torchx_opts()},
+      {:nx, "~> 1.0", nx_opts()},
+      {:exla, "~> 1.0", [only: :test] ++ exla_opts()},
+      {:torchx, "~> 1.0", [only: :test] ++ torchx_opts()},
       {:ex_doc, "~> 0.34", only: :docs},
       {:table_rex, "~> 3.1 or ~> 4.1", optional: true},
       {:kino, "~> 0.7", optional: true},
@@ -62,21 +62,20 @@ defmodule Axon.MixProject do
     ]
   end
 
-  # TODO: back to the released Nx once buffer donation (`Nx.donatable/1`,
-  # used by `Axon.Loop.run/4`'s `:donate_state?`) is in a release.
+  # Nx stays an override because kino still declares `nx ~> 0.1`
   defp nx_opts do
     if path = System.get_env("AXON_NX_PATH") do
       [path: path, override: true]
     else
-      [github: "elixir-nx/nx", sparse: "nx", override: true]
+      [override: true]
     end
   end
 
   defp exla_opts do
     if path = System.get_env("AXON_EXLA_PATH") do
-      [path: path]
+      [path: path, override: true]
     else
-      [github: "elixir-nx/nx", sparse: "exla", override: true]
+      []
     end
   end
 

mix.lock

@@ -1,10 +1,10 @@
 %{
   "castore": {:hex, :castore, "1.0.16", "8a4f9a7c8b81cda88231a08fe69e3254f16833053b23fa63274b05cbc61d2a1e", [:mix], [], "hexpm", "33689203a0eaaf02fcd0e86eadfbcf1bd636100455350592e7e2628564022aaf"},
-  "complex": {:hex, :complex, "0.7.0", "695632ef9487517aa5d57edd1697801079d622414cb2e1a7cf538b1f9a50f205", [:mix], [], "hexpm", "0ee39c0803129f546e7f3f640da8f021c9e659402bf59da6f7f2c4848f068f8d"},
+  "complex": {:hex, :complex, "1.0.0", "32075df521760e83d808652f52dd57f21712d47d5918a328709da78358387507", [:mix], [], "hexpm", "b4fd7ab305712472521043aa4322f7e4cdd85ca3ddf79c9c2c338d6663adc8d8"},
   "earmark_parser": {:hex, :earmark_parser, "1.4.41", "ab34711c9dc6212dda44fcd20ecb87ac3f3fce6f0ca2f28d4a00e4154f8cd599", [:mix], [], "hexpm", "a81a04c7e34b6617c2792e291b5a2e57ab316365c2644ddc553bb9ed863ebefa"},
   "elixir_make": {:hex, :elixir_make, "0.10.0", "16577e2583a79bb79237bbff349619ef5d80afffc07eac6e4faf0d00e2ddaf7d", [:mix], [], "hexpm", "dc1f09fb7fa68866b886abd5f0f3c83553b1a19a52359a899e92af1bb3b31982"},
   "ex_doc": {:hex, :ex_doc, "0.34.2", "13eedf3844ccdce25cfd837b99bea9ad92c4e511233199440488d217c92571e8", [:mix], [{:earmark_parser, "~> 1.4.39", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "5ce5f16b41208a50106afed3de6a2ed34f4acfd65715b82a0b84b49d995f95c1"},
-  "exla": {:git, "https://github.com/elixir-nx/nx.git", "5ad4ec454d48e2b34820016bcf508ee04a134d56", [sparse: "exla"]},
+  "exla": {:hex, :exla, "1.0.0", "1c24b1e38416c954bf45851e7591e945876b76ac37e6ba992f01bb72ea615163", [:make, :mix], [{:elixir_make, "~> 0.6", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:fine, "~> 0.1", [hex: :fine, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.0", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:nx, "~> 1.0", [hex: :nx, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}, {:xla, "~> 0.10.0", [hex: :xla, repo: "hexpm", optional: false]}], "hexpm", "f89b030d368372fe62aacecc773e2d096d812d1afe55acfed333ac9a43864ba7"},
   "fine": {:hex, :fine, "0.1.6", "4bf7151493443c454aac9f2fa2f34f5fefd0346a83fb5586a016c4a135c63247", [:mix], [], "hexpm", "5638eb4495488e885ebec167fa57973e5c35e1a50c344eb7666c90ec1c4e3b12"},
   "fss": {:hex, :fss, "0.1.1", "9db2344dbbb5d555ce442ac7c2f82dd975b605b50d169314a20f08ed21e08642", [:mix], [], "hexpm", "78ad5955c7919c3764065b21144913df7515d52e228c09427a004afe9c1a16b0"},
   "kino": {:hex, :kino, "0.14.1", "c499afb1cd0be462feaf0a75c0631aa65aacc545b1c10f431b439b74f104be22", [:mix], [{:fss, "~> 0.1.0", [hex: :fss, repo: "hexpm", optional: false]}, {:nx, "~> 0.1", [hex: :nx, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}, {:table, "~> 0.1.2", [hex: :table, repo: "hexpm", optional: false]}], "hexpm", "090aea1aaa267e42e5ac24ee6bc5ed515aecc0a9edb8619aa4ee839201e704aa"},
@@ -15,13 +15,13 @@
   "makeup_syntect": {:hex, :makeup_syntect, "0.1.3", "ae2c3437f479ea50d08d794acaf02a2f3a8c338dd1f757f6b237c42eb27fcde1", [:mix], [{:makeup, "~> 1.2", [hex: :makeup, repo: "hexpm", optional: false]}, {:rustler, "~> 0.36.1", [hex: :rustler, repo: "hexpm", optional: true]}, {:rustler_precompiled, "~> 0.8.2", [hex: :rustler_precompiled, repo: "hexpm", optional: false]}], "hexpm", "a27bd3bd8f7b87465d110295a33ed1022202bea78701bd2bbeadfb45d690cdbf"},
   "nimble_parsec": {:hex, :nimble_parsec, "1.4.2", "8efba0122db06df95bfaa78f791344a89352ba04baedd3849593bfce4d0dc1c6", [:mix], [], "hexpm", "4b21398942dda052b403bbe1da991ccd03a053668d147d53fb8c4e0efe09c973"},
   "nimble_pool": {:hex, :nimble_pool, "1.1.0", "bf9c29fbdcba3564a8b800d1eeb5a3c58f36e1e11d7b7fb2e084a643f645f06b", [:mix], [], "hexpm", "af2e4e6b34197db81f7aad230c1118eac993acc0dae6bc83bac0126d4ae0813a"},
-  "nx": {:git, "https://github.com/elixir-nx/nx.git", "5ad4ec454d48e2b34820016bcf508ee04a134d56", [sparse: "nx"]},
+  "nx": {:hex, :nx, "1.0.0", "83d586079cb37d11a4cf8c0edaf2eccafe88aa2369e35ebdb1c3d78c43c8eb33", [:mix], [{:complex, "~> 1.0", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "e06c7587efff7a3dfc1e2c92705943e319db20070359e9850c3921aaf26d6998"},
   "polaris": {:hex, :polaris, "0.1.0", "dca61b18e3e801ecdae6ac9f0eca5f19792b44a5cb4b8d63db50fc40fc038d22", [:mix], [{:nx, "~> 0.5", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "13ef2b166650e533cb24b10e2f3b8ab4f2f449ba4d63156e8c569527f206e2c2"},
   "rustler_precompiled": {:hex, :rustler_precompiled, "0.8.2", "5f25cbe220a8fac3e7ad62e6f950fcdca5a5a5f8501835d2823e8c74bf4268d5", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: false]}, {:rustler, "~> 0.23", [hex: :rustler, repo: "hexpm", optional: true]}], "hexpm", "63d1bd5f8e23096d1ff851839923162096364bac8656a4a3c00d1fff8e83ee0a"},
   "table": {:hex, :table, "0.1.2", "87ad1125f5b70c5dea0307aa633194083eb5182ec537efc94e96af08937e14a8", [:mix], [], "hexpm", "7e99bc7efef806315c7e65640724bf165c3061cdc5d854060f74468367065029"},
   "table_rex": {:hex, :table_rex, "4.1.0", "fbaa8b1ce154c9772012bf445bfb86b587430fb96f3b12022d3f35ee4a68c918", [:mix], [], "hexpm", "95932701df195d43bc2d1c6531178fc8338aa8f38c80f098504d529c43bc2601"},
   "telemetry": {:hex, :telemetry, "1.4.2", "a0cb522801dffb1c49fe6e30561badffc7b6d0e180db1300df759faa22062855", [:rebar3], [], "hexpm", "928f6495066506077862c0d1646609eed891a4326bee3126ba54b60af61febb1"},
-  "torchx": {:hex, :torchx, "0.13.0", "695b84e0565612daadb4c3eda784aa10d5a66f6046e5b146a8377c4aa26bd48c", [:mix], [{:fine, "~> 0.1.0", [hex: :fine, repo: "hexpm", optional: false]}, {:nx, "~> 0.13.0", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "2c3b7a2cb1639f0a010f15334f9b6a99f0c4fa60d8274a3be2b4614c06d627de"},
+  "torchx": {:hex, :torchx, "1.0.0", "5945d22cc33158061ad6102b1be437fc3ff34c785e485e4eabb15c87d07202a0", [:mix], [{:fine, "~> 0.1.0", [hex: :fine, repo: "hexpm", optional: false]}, {:nx, "~> 1.0", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "288790ac18f7d585e3989d178cb0eedbef7bfa3f8ce8fab375e679d40190b0d3"},
   "vega_lite": {:hex, :vega_lite, "0.1.9", "d7a288665f916181b68d0a3617f1b3611d16a4dcd5fafb51b847b71db1159d4c", [:mix], [{:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: false]}], "hexpm", "c6a056e763162198e73ae6dfb46c09753bb0298474410fd085074e1cdcee7418"},
   "xla": {:hex, :xla, "0.10.0", "41121e9f011456242d3a79b9289910ce43419be0b0e7ebe67cc1292c6b3f232f", [:make, :mix], [{:elixir_make, "~> 0.4", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "f57d91aea6e661b52bf12239316c598679e9170628122bbd941235f040122bc6"},
 }

test/axon/loop_test.exs

@@ -663,7 +663,7 @@ defmodule Axon.LoopTest do
                      |> Loop.run(data, Axon.ModelState.empty())
                    end
     end
- 
+
     test "logs epoch and iteration" do
       model = Axon.input("input", shape: {nil, 1}) |> Axon.dense(1)
       data = [{Nx.tensor([[1.0]]), Nx.tensor([[2.0]])}]