diff --git a/Cargo.lock b/Cargo.lock
index 53c7626..275740a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -35,9 +35,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
 dependencies = [
  "cfg-if",
- "getrandom 0.3.4",
  "once_cell",
- "serde",
  "version_check",
  "zerocopy",
 ]
@@ -147,6 +145,8 @@ dependencies = [
  "event-listener-strategy",
  "futures-core",
  "pin-project-lite",
+ "portable-atomic",
+ "portable-atomic-util",
 ]
 
 [[package]]
@@ -234,12 +234,6 @@ dependencies = [
  "windows-link",
 ]
 
-[[package]]
-name = "base64"
-version = "0.13.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8"
-
 [[package]]
 name = "base64"
 version = "0.22.1"
@@ -252,36 +246,6 @@ version = "1.8.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06"
 
-[[package]]
-name = "bincode"
-version = "2.0.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "36eaf5d7b090263e8150820482d5d93cd964a81e4019913c972f4edcc6edb740"
-dependencies = [
- "serde",
- "unty",
-]
-
-[[package]]
-name = "bindgen"
-version = "0.71.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3"
-dependencies = [
- "bitflags 2.11.1",
- "cexpr",
- "clang-sys",
- "itertools 0.13.0",
- "log",
- "prettyplease",
- "proc-macro2",
- "quote",
- "regex",
- "rustc-hash 2.1.2",
- "shlex",
- "syn 2.0.117",
-]
-
 [[package]]
 name = "bit-set"
 version = "0.5.3"
@@ -300,6 +264,15 @@ dependencies = [
  "bit-vec 0.9.1",
 ]
 
+[[package]]
+name = "bit-set"
+version = "0.10.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09ec2f926cc3060f09db9ebc5b52823d85268d24bb917e472c0c4bea35780a7d"
+dependencies = [
+ "bit-vec 0.9.1",
+]
+
 [[package]]
 name = "bit-vec"
 version = "0.6.3"
@@ -399,26 +372,19 @@ checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649"
 
 [[package]]
 name = "burn"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "39474be398d30e08681f1f6a8ff446b1e4f1403b5cf7749649a8871fd5945f3a"
+checksum = "d450b8a4aefca19ce1dd51ee6980eaf379b6d238ced4de37535e75f72be7a532"
 dependencies = [
- "burn-autodiff",
- "burn-candle",
+ "burn-autodiff 0.22.0-pre.2",
  "burn-core",
- "burn-cpu",
  "burn-cuda",
- "burn-dispatch",
- "burn-flex",
- "burn-ir",
- "burn-ndarray",
+ "burn-ir 0.22.0-pre.2",
+ "burn-ndarray 0.22.0-pre.2",
  "burn-nn",
  "burn-optim",
- "burn-rocm",
- "burn-router",
- "burn-std",
+ "burn-std 0.22.0-pre.2",
  "burn-store",
- "burn-tch",
  "burn-train",
  "burn-vision",
  "burn-wgpu",
@@ -430,14 +396,31 @@ version = "0.21.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3b93a80e43bfab909399444b29ce3894ff51f7e879720bfc75b6007ae6a01c3d"
 dependencies = [
- "burn-backend",
- "burn-std",
+ "burn-backend 0.21.0",
+ "burn-std 0.21.0",
+ "derive-new",
+ "hashbrown 0.16.1",
+ "log",
+ "num-traits",
+ "portable-atomic",
+ "spin",
+]
+
+[[package]]
+name = "burn-autodiff"
+version = "0.22.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7c8a55150fdff1d1e8e1aca98010fb10f443e5808cdf1bad6cf7d0b096f1d7fe"
+dependencies = [
+ "burn-backend 0.22.0-pre.2",
+ "burn-std 0.22.0-pre.2",
  "derive-new",
  "hashbrown 0.16.1",
  "log",
  "num-traits",
  "parking_lot",
  "portable-atomic",
+ "portable-atomic-util",
  "spin",
 ]
 
@@ -447,48 +430,69 @@ version = "0.21.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "64491519ac0867f550fa9c726cd443e5db94608c629050986cf2614c8c5ad39f"
 dependencies = [
- "burn-std",
+ "burn-std 0.21.0",
  "bytemuck",
- "cubecl",
+ "cubecl 0.10.0",
  "derive-new",
  "enumset",
  "hashbrown 0.16.1",
  "num-traits",
  "portable-atomic-util",
  "rand 0.10.1",
- "rand_distr 0.6.0",
+ "rand_distr",
  "serde",
  "spin",
  "thiserror 2.0.18",
 ]
 
 [[package]]
-name = "burn-candle"
-version = "0.21.0"
+name = "burn-backend"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "917fbbe1238d80ec1483c2c360ed9fa0146fdf57b2e66015c0824ccf5155276c"
+checksum = "b8157cd9dfe164c4a3e00183a4fbaa26789209ecf13ba55906185b532a20f1c5"
 dependencies = [
- "burn-backend",
- "burn-std",
- "candle-core",
+ "burn-std 0.22.0-pre.2",
+ "bytemuck",
+ "cubecl 0.11.0-pre.2",
  "derive-new",
+ "enumset",
+ "hashbrown 0.16.1",
+ "num-traits",
+ "oneshot",
+ "portable-atomic-util",
+ "rand 0.10.1",
+ "rand_distr",
+ "serde",
+ "spin",
+]
+
+[[package]]
+name = "burn-backend-extension"
+version = "0.22.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4fdfe800838c9949eecd9b0f5250362148baf1f2b1dce710ac3775cf84fba1a3"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 2.0.117",
 ]
 
 [[package]]
 name = "burn-core"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5bde1e3b8e7e39b0aa71cd5905207ceea5c3478bec47bdb64acf72dcd0d6f978"
+checksum = "7861d5f03dfa489c05313cfd69bb02be59f4cece8cdc54731e0bb248abd2bb00"
 dependencies = [
  "ahash",
- "bincode",
+ "burn-backend 0.22.0-pre.2",
+ "burn-backend-extension",
  "burn-dataset",
  "burn-derive",
- "burn-std",
+ "burn-dispatch",
+ "burn-pack",
+ "burn-std 0.22.0-pre.2",
  "burn-tensor",
- "data-encoding",
  "derive-new",
- "flate2",
  "half",
  "hashbrown 0.16.1",
  "log",
@@ -497,37 +501,35 @@ dependencies = [
  "portable-atomic-util",
  "rand 0.10.1",
  "regex",
- "rmp-serde",
  "serde",
  "serde_json",
  "spin",
- "uuid",
 ]
 
 [[package]]
 name = "burn-cpu"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "681478c5438ab6c753a55b85f5b42e5bd4c1f69a8564f2afa2f698ac396ba556"
+checksum = "357ce26165dcaf27c8acfdff97170a6ce724db3c4076252fbb0669ae274fec53"
 dependencies = [
- "burn-backend",
+ "burn-backend 0.22.0-pre.2",
  "burn-cubecl",
  "burn-fusion",
- "cubecl",
+ "cubecl 0.11.0-pre.2",
 ]
 
 [[package]]
 name = "burn-cubecl"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "26efbec96ca3f691593c966fa90f98ef8a72867f62627904920b619ba3a9e6b6"
+checksum = "3a1db3d0cd0d081e2dd0b8dce41847bc88495cc3e9e89c47c391f8c6940263bd"
 dependencies = [
- "burn-backend",
+ "burn-backend 0.22.0-pre.2",
  "burn-cubecl-fusion",
  "burn-fusion",
- "burn-ir",
- "burn-std",
- "cubecl",
+ "burn-ir 0.22.0-pre.2",
+ "burn-std 0.22.0-pre.2",
+ "cubecl 0.11.0-pre.2",
  "cubek",
  "derive-new",
  "futures-lite",
@@ -538,39 +540,40 @@ dependencies = [
 
 [[package]]
 name = "burn-cubecl-fusion"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6993d15b8ed588f55626e835ca768c1c3c7aa7a79d59beff4c987c40fc82d1f4"
+checksum = "40c80f796e4c88fd196022ca1f12a5d04ae28f175cb8be82ca364446eb7b0e7f"
 dependencies = [
- "burn-backend",
+ "burn-backend 0.22.0-pre.2",
  "burn-fusion",
- "burn-ir",
- "burn-std",
- "cubecl",
+ "burn-ir 0.22.0-pre.2",
+ "burn-std 0.22.0-pre.2",
+ "cubecl 0.11.0-pre.2",
  "cubek",
  "derive-new",
  "hashbrown 0.16.1",
+ "rmp-serde",
  "serde",
  "spin",
 ]
 
 [[package]]
 name = "burn-cuda"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a6c66b49136fc11f59773054358f401fbec7fa0d69615f55ef22c29cf241c9b7"
+checksum = "3d2f7d0d2dfa0221d87e095944807f38fb201582b9b9c62d72f47019a5284b11"
 dependencies = [
- "burn-backend",
+ "burn-backend 0.22.0-pre.2",
  "burn-cubecl",
  "burn-fusion",
- "cubecl",
+ "cubecl 0.11.0-pre.2",
 ]
 
 [[package]]
 name = "burn-dataset"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ba469fba38357c8bb65458873ddfd82aa97534aec0b3bae93d4c44604fc63839"
+checksum = "90aecfb8fa00b8e04053171972c64d0b22f4eaa948363665443bb24cae93fd3d"
 dependencies = [
  "csv",
  "derive-new",
@@ -587,9 +590,9 @@ dependencies = [
 
 [[package]]
 name = "burn-derive"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dce7be43cdfc9903c43dbaf3f4821543e0497d79047d8fc9ee9084e448844446"
+checksum = "b0dc2ecbef19c9f73f6e567715f20968ac459490a137c5c6f30d6c95d64a3172"
 dependencies = [
  "derive-new",
  "proc-macro2",
@@ -599,16 +602,17 @@ dependencies = [
 
 [[package]]
 name = "burn-dispatch"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c2087b9e12323e0d25cc0ebdccfee40427dfdec4a23fa6d3f49ba2aa1f94ac17"
+checksum = "b5bf565da978d065420c0b750417e5bc6ee72e5a168e0b2b15fc300439a6d622"
 dependencies = [
- "burn-autodiff",
- "burn-backend",
+ "burn-autodiff 0.22.0-pre.2",
+ "burn-backend 0.22.0-pre.2",
  "burn-cpu",
  "burn-cuda",
  "burn-flex",
- "burn-ndarray",
+ "burn-ndarray 0.22.0-pre.2",
+ "burn-remote",
  "burn-rocm",
  "burn-tch",
  "burn-wgpu",
@@ -617,32 +621,35 @@ dependencies = [
 
 [[package]]
 name = "burn-flex"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0098bf6555c151517b4c98ca0e8ed1ab6a976e40f5f4e967c117b00eee21e192"
+checksum = "3af60c5cda3823448c723d6e3840223e806bb122a639689421a2ca1164decbaa"
 dependencies = [
  "aligned-vec",
- "burn-backend",
- "burn-ir",
- "burn-std",
+ "burn-backend 0.22.0-pre.2",
+ "burn-ir 0.22.0-pre.2",
+ "burn-std 0.22.0-pre.2",
  "bytemuck",
  "gemm",
  "half",
  "libm",
  "macerator",
  "num-traits",
+ "once_cell",
+ "portable-atomic",
+ "portable-atomic-util",
  "rayon",
 ]
 
 [[package]]
 name = "burn-fusion"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "95e58fcfd52a60cdf2f7da34fbebf0ed38f960fbd8675d384bacd5795cf2bdf2"
+checksum = "0752c69048173820fd19cdbcb11f50e780dc4aaa2fee9f68ab75a6e0365b975e"
 dependencies = [
- "burn-backend",
- "burn-ir",
- "burn-std",
+ "burn-backend 0.22.0-pre.2",
+ "burn-ir 0.22.0-pre.2",
+ "burn-std 0.22.0-pre.2",
  "derive-new",
  "hashbrown 0.16.1",
  "log",
@@ -658,7 +665,18 @@ version = "0.21.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a81a55c32f35c2265d9e15ba2e796013d9780d5a49f9e21ea0ac4d29609dbf13"
 dependencies = [
- "burn-backend",
+ "burn-backend 0.21.0",
+ "hashbrown 0.16.1",
+ "serde",
+]
+
+[[package]]
+name = "burn-ir"
+version = "0.22.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26d4ac0a2fcd725d7d429589944f2b894bc1ff5da01d40598f46b9561fc8b717"
+dependencies = [
+ "burn-backend 0.22.0-pre.2",
  "hashbrown 0.16.1",
  "serde",
 ]
@@ -670,13 +688,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "dceb9692e292782bab7ad0259e8e8f5ee80ba2b0329a78f0ba589a26a9633dd6"
 dependencies = [
  "atomic_float",
- "burn-autodiff",
- "burn-backend",
- "burn-ir",
- "burn-std",
+ "burn-autodiff 0.21.0",
+ "burn-backend 0.21.0",
+ "burn-ir 0.21.0",
+ "burn-std 0.21.0",
  "bytemuck",
  "const-random",
- "itertools 0.14.0",
+ "itertools",
  "libm",
  "macerator",
  "matrixmultiply",
@@ -691,10 +709,32 @@ dependencies = [
 ]
 
 [[package]]
-name = "burn-nn"
-version = "0.21.0"
+name = "burn-ndarray"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "93e4acd90806fa8663610f1071f3a8992eb98637d2d0816ee3062b334f61af00"
+checksum = "5e59e56039e0ddd460741b95ccc5cb275de365c3127927a767b16421f0b4e157"
+dependencies = [
+ "atomic_float",
+ "burn-backend 0.22.0-pre.2",
+ "burn-ir 0.22.0-pre.2",
+ "burn-std 0.22.0-pre.2",
+ "const-random",
+ "libm",
+ "macerator",
+ "matrixmultiply",
+ "ndarray 0.17.2",
+ "num-traits",
+ "paste",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rand 0.10.1",
+]
+
+[[package]]
+name = "burn-nn"
+version = "0.22.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3b5d3b52e279419f3ea529bd549ae0d593ae4514b795df1a157b3f582c1cde1c"
 dependencies = [
  "burn-core",
  "num-traits",
@@ -702,11 +742,13 @@ dependencies = [
 
 [[package]]
 name = "burn-optim"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4c5ea466ae0b85bf18f7656cdc152050dc8c581057fbbe33bbed267e22600228"
+checksum = "b65488c481e479158aa85950e9a128d42548d920f89b503adafec86c12609311"
 dependencies = [
  "burn-core",
+ "burn-derive",
+ "burn-pack",
  "derive-new",
  "hashbrown 0.16.1",
  "log",
@@ -715,28 +757,65 @@ dependencies = [
 ]
 
 [[package]]
-name = "burn-rocm"
-version = "0.21.0"
+name = "burn-pack"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "66af82fdaaf2ad0fa4b0dae90119aa94f333105d76d520e0eb48917717c06fde"
+checksum = "93325a6545a605b5079d77957970b26f920672c12555589b34ae67242e3f4717"
 dependencies = [
- "burn-backend",
+ "burn-std 0.22.0-pre.2",
+ "byteorder",
+ "ciborium",
+ "serde",
+]
+
+[[package]]
+name = "burn-remote"
+version = "0.22.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fa8db206a100d838eeb100b8525ff8e3051958ee822fcfecd736fe58eddaeeaa"
+dependencies = [
+ "burn-backend 0.22.0-pre.2",
+ "burn-fusion",
+ "burn-ir 0.22.0-pre.2",
+ "burn-router",
+ "burn-std 0.22.0-pre.2",
+ "bytes",
+ "futures-util",
+ "gloo-timers",
+ "log",
+ "rand 0.10.1",
+ "rmp-serde",
+ "serde",
+ "serde_bytes",
+ "tokio",
+ "wasm-bindgen-futures",
+]
+
+[[package]]
+name = "burn-rocm"
+version = "0.22.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8667b9c5003d82f2c0511a8a8aef0bddf68874f261d549653a3e7e0fabe37164"
+dependencies = [
+ "burn-backend 0.22.0-pre.2",
  "burn-cubecl",
  "burn-fusion",
- "cubecl",
+ "cubecl 0.11.0-pre.2",
 ]
 
 [[package]]
 name = "burn-router"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7960a5d43918c3158629da2bd7956e38dda16ef54de95cca83b4b7b1477e6185"
+checksum = "91cc6e085ed958e988840a4fd3c4ec684a0479dcd3089bcf1b55babd8735eec7"
 dependencies = [
- "burn-backend",
- "burn-ir",
- "burn-std",
+ "burn-backend 0.22.0-pre.2",
+ "burn-fusion",
+ "burn-ir 0.22.0-pre.2",
+ "burn-std 0.22.0-pre.2",
  "hashbrown 0.16.1",
  "log",
+ "serde",
  "spin",
 ]
 
@@ -748,9 +827,8 @@ checksum = "81c7eb60eb2c316854af5b214fd55f59fc7b0e25dfde7db671f09cc49f269048"
 dependencies = [
  "bytemuck",
  "bytes",
- "cubecl",
- "cubecl-common",
- "cubecl-zspace",
+ "cubecl-common 0.10.0",
+ "cubecl-zspace 0.10.0",
  "half",
  "num-traits",
  "serde",
@@ -759,13 +837,41 @@ dependencies = [
 ]
 
 [[package]]
-name = "burn-store"
-version = "0.21.0"
+name = "burn-std"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2f64a4fe2a0616d18f07d1d0d19c79e655e70e1828cbcb5e9df7b52211091023"
+checksum = "7327d31c17e6eda647a7e43752353d5481d4cdca6c89112cbcd6e5402bb45264"
+dependencies = [
+ "ahash",
+ "bytemuck",
+ "bytes",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-zspace 0.11.0-pre.2",
+ "data-encoding",
+ "derive-new",
+ "enumset",
+ "half",
+ "hashbrown 0.16.1",
+ "num-traits",
+ "portable-atomic-util",
+ "rand 0.10.1",
+ "rand_distr",
+ "serde",
+ "smallvec",
+ "spin",
+ "thiserror 2.0.18",
+ "uuid",
+]
+
+[[package]]
+name = "burn-store"
+version = "0.22.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ba01d3f12054d79d2f715c387ca62836d3a18a668f6f5838044c4871f38b6825"
 dependencies = [
  "burn-core",
- "burn-tensor",
+ "burn-pack",
  "byteorder",
  "bytes",
  "half",
@@ -778,11 +884,11 @@ dependencies = [
 
 [[package]]
 name = "burn-tch"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dfa36b630bd7a790395b4561f7d0e315c0d5e1c557db7dd4cf1b48d5c0d4aff2"
+checksum = "5b5ff862775a74d15b7e11205368e9cc21e3d79da62be91271efdc90a62e737d"
 dependencies = [
- "burn-backend",
+ "burn-backend 0.22.0-pre.2",
  "cc",
  "libc",
  "log",
@@ -792,12 +898,13 @@ dependencies = [
 
 [[package]]
 name = "burn-tensor"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "223ed3804bb9436e401fc57b87e44c86267070b870813520ed9f706c80c81442"
+checksum = "85e85e6beb63d3ccd18fa5c1f71e8e3806a7f940cec8769f522b9d8cd173c999"
 dependencies = [
- "burn-backend",
- "burn-std",
+ "burn-backend 0.22.0-pre.2",
+ "burn-dispatch",
+ "burn-std 0.22.0-pre.2",
  "colored",
  "derive-new",
  "num-traits",
@@ -806,14 +913,15 @@ dependencies = [
 
 [[package]]
 name = "burn-train"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fa65e2e569e571bb8dc9d37aa8fcf715f522f8159c6ee87f665c1961fee5259f"
+checksum = "da08a29fff342175c922437a36afd296304abbf4defb0620e3d177406dc42d6b"
 dependencies = [
  "async-channel",
  "burn-core",
  "burn-flex",
  "burn-optim",
+ "burn-std 0.22.0-pre.2",
  "derive-new",
  "log",
  "rand 0.10.1",
@@ -824,22 +932,23 @@ dependencies = [
  "tracing-appender",
  "tracing-core",
  "tracing-subscriber",
+ "unicode-width",
 ]
 
 [[package]]
 name = "burn-vision"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38b5a3e9096977462cd22bf713087f9777e427f7693508871f86f5b1e8347a70"
+checksum = "550d4ecb73cc1ecf1aa3160e9844192bba76c6969a7c2da1ba885c4596278d3f"
 dependencies = [
  "aligned-vec",
  "bon",
+ "burn-core",
  "burn-cubecl",
  "burn-fusion",
- "burn-ir",
- "burn-tensor",
+ "burn-ir 0.22.0-pre.2",
  "bytemuck",
- "cubecl",
+ "cubecl 0.11.0-pre.2",
  "derive-new",
  "half",
  "image",
@@ -852,14 +961,14 @@ dependencies = [
 
 [[package]]
 name = "burn-wgpu"
-version = "0.21.0"
+version = "0.22.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6eb009254af15922cb37814f3b3b61043046193ba2fde97c3879a25fbd51a92e"
+checksum = "798e66d9efc792e6e646e16d6f531a2c6e503bbe92f2b017c84f2306fcc9e2d3"
 dependencies = [
- "burn-backend",
+ "burn-backend 0.22.0-pre.2",
  "burn-cubecl",
  "burn-fusion",
- "cubecl",
+ "cubecl 0.11.0-pre.2",
 ]
 
 [[package]]
@@ -923,30 +1032,6 @@ dependencies = [
  "pkg-config",
 ]
 
-[[package]]
-name = "candle-core"
-version = "0.10.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6bd9895436c1ba5dc1037a19935d084b838db066ff4e15ef7dded020b7c12a4a"
-dependencies = [
- "byteorder",
- "float8",
- "gemm",
- "half",
- "libm",
- "memmap2",
- "num-traits",
- "num_cpus",
- "rand 0.9.4",
- "rand_distr 0.5.1",
- "rayon",
- "safetensors 0.7.0",
- "thiserror 2.0.18",
- "tokenizers",
- "yoke",
- "zip 7.2.0",
-]
-
 [[package]]
 name = "caseless"
 version = "0.2.2"
@@ -977,15 +1062,6 @@ dependencies = [
  "shlex",
 ]
 
-[[package]]
-name = "cexpr"
-version = "0.6.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6fac387a98bb7c37292057cffc56d62ecb629900026402633ae9160df93a8766"
-dependencies = [
- "nom 7.1.3",
-]
-
 [[package]]
 name = "cfg-if"
 version = "1.0.4"
@@ -1058,17 +1134,6 @@ dependencies = [
  "inout",
 ]
 
-[[package]]
-name = "clang-sys"
-version = "1.8.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0b023947811758c97c59bf9d1c188fd619ad4718dcaa767947df1cadb14f39f4"
-dependencies = [
- "glob",
- "libc",
- "libloading 0.8.9",
-]
-
 [[package]]
 name = "cli"
 version = "0.1.0"
@@ -1113,7 +1178,6 @@ dependencies = [
  "itoa",
  "rustversion",
  "ryu",
- "serde",
  "static_assertions",
 ]
 
@@ -1138,6 +1202,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973"
 dependencies = [
  "crossbeam-utils",
+ "portable-atomic",
 ]
 
 [[package]]
@@ -1348,14 +1413,29 @@ version = "0.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fd203fef6e359472e4cb8f6c0ef3eca062815a1edd560eb6d6c1d5c1397838d9"
 dependencies = [
- "cubecl-core",
+ "cubecl-core 0.10.0",
+ "cubecl-cuda 0.10.0",
+ "cubecl-ir 0.10.0",
+ "cubecl-runtime 0.10.0",
+ "cubecl-wgpu 0.10.0",
+ "half",
+]
+
+[[package]]
+name = "cubecl"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2029987fbb8d073eabc5ca09887846c2db6928fa81f90afc6b883ce819425884"
+dependencies = [
+ "cubecl-core 0.11.0-pre.2",
  "cubecl-cpu",
- "cubecl-cuda",
+ "cubecl-cuda 0.11.0-pre.2",
+ "cubecl-environment",
  "cubecl-hip",
- "cubecl-ir",
- "cubecl-runtime",
+ "cubecl-ir 0.11.0-pre.2",
+ "cubecl-runtime 0.11.0-pre.2",
  "cubecl-std",
- "cubecl-wgpu",
+ "cubecl-wgpu 0.11.0-pre.2",
  "half",
 ]
 
@@ -1394,13 +1474,37 @@ dependencies = [
  "serde_json",
  "spin",
  "toml",
- "tracing",
  "tynm",
  "wasm-bindgen-futures",
  "web-time",
  "xxhash-rust",
 ]
 
+[[package]]
+name = "cubecl-common"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c931f1ac56ee174030fe64374b8bdf5638b9e8376f885fc8487f75debb236ae5"
+dependencies = [
+ "bytemuck",
+ "cfg_aliases",
+ "cubecl-environment",
+ "derive-new",
+ "derive_more",
+ "float4",
+ "float8",
+ "half",
+ "hashbrown 0.17.1",
+ "log",
+ "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
+ "serde",
+ "spin",
+ "tynm",
+ "xxhash-rust",
+]
+
 [[package]]
 name = "cubecl-core"
 version = "0.10.0"
@@ -1409,11 +1513,11 @@ checksum = "a7522e7acf25d7848032c7b5270780f9f2abe84e13c7ed6345aa27ba70c312ca"
 dependencies = [
  "bitflags 2.11.1",
  "bytemuck",
- "cubecl-common",
- "cubecl-ir",
- "cubecl-macros",
- "cubecl-runtime",
- "cubecl-zspace",
+ "cubecl-common 0.10.0",
+ "cubecl-ir 0.10.0",
+ "cubecl-macros 0.10.0",
+ "cubecl-runtime 0.10.0",
+ "cubecl-zspace 0.10.0",
  "derive-new",
  "derive_more",
  "enumset",
@@ -1425,8 +1529,38 @@ dependencies = [
  "paste",
  "serde",
  "serde_json",
+ "variadics_please 1.1.0",
+]
+
+[[package]]
+name = "cubecl-core"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d6fd5fb5fbf719e1a86df673dddae9e5183b08f0744308b2558fcc0632f5fe9"
+dependencies = [
+ "bitflags 2.11.1",
+ "bytemuck",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-ir 0.11.0-pre.2",
+ "cubecl-macros 0.11.0-pre.2",
+ "cubecl-runtime 0.11.0-pre.2",
+ "cubecl-zspace 0.11.0-pre.2",
+ "derive-new",
+ "derive_more",
+ "enumset",
+ "float-ord",
+ "half",
+ "hashbrown 0.17.1",
+ "itertools",
+ "log",
+ "num-integer",
+ "num-traits",
+ "paste",
+ "serde",
+ "serde_json",
  "tracing",
- "variadics_please",
+ "variadics_please 2.0.0",
 ]
 
 [[package]]
@@ -1436,36 +1570,60 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "411af04828cbf4583ecd388579fb30cd7a644eec19a334bb8f0d9d08522e1458"
 dependencies = [
  "bytemuck",
- "cubecl-common",
- "cubecl-core",
- "cubecl-opt",
- "cubecl-runtime",
+ "cubecl-common 0.10.0",
+ "cubecl-core 0.10.0",
+ "cubecl-opt 0.10.0",
+ "cubecl-runtime 0.10.0",
  "derive-new",
  "half",
- "itertools 0.14.0",
+ "itertools",
+ "log",
+]
+
+[[package]]
+name = "cubecl-cpp"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1f43128c27f790ff22012679945c2b8972296979c52e5c323b54be830069e6e9"
+dependencies = [
+ "bytemuck",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-core 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-opt 0.11.0-pre.2",
+ "cubecl-runtime 0.11.0-pre.2",
+ "derive-new",
+ "half",
+ "itertools",
  "log",
 ]
 
 [[package]]
 name = "cubecl-cpu"
-version = "0.10.0"
+version = "0.11.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f572143f54e42a49ee0635a4ec36005f639e62be029e4f49890c808985981b35"
+checksum = "8274cac78ff1528b075eb0240b576690c419e13c4621dc1c596ea0a6fce809ca"
 dependencies = [
  "bytemuck",
- "cubecl-common",
- "cubecl-core",
- "cubecl-opt",
- "cubecl-runtime",
+ "crossbeam-utils",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-core 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-opt 0.11.0-pre.2",
+ "cubecl-runtime 0.11.0-pre.2",
  "cubecl-std",
  "derive-new",
  "half",
+ "libc",
  "log",
  "paste",
  "serde",
+ "smallvec",
+ "spin",
  "sysinfo",
  "tracel-llvm",
  "tracel-llvm-bundler",
+ "winapi",
 ]
 
 [[package]]
@@ -1475,10 +1633,29 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b6b0a69ff45688d322ad8e92c8bf645167b9ca490fa8fa087fc6adac8c5e46be"
 dependencies = [
  "bytemuck",
- "cubecl-common",
- "cubecl-core",
- "cubecl-cpp",
- "cubecl-runtime",
+ "cubecl-common 0.10.0",
+ "cubecl-core 0.10.0",
+ "cubecl-cpp 0.10.0",
+ "cubecl-runtime 0.10.0",
+ "cudarc",
+ "derive-new",
+ "half",
+ "log",
+ "serde",
+]
+
+[[package]]
+name = "cubecl-cuda"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3e51bbfb4834e79972d7248bb8f69a747e848415f6e06e01aa8529cc8a29ad69"
+dependencies = [
+ "bytemuck",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-core 0.11.0-pre.2",
+ "cubecl-cpp 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-runtime 0.11.0-pre.2",
  "cudarc",
  "derive-new",
  "half",
@@ -1488,17 +1665,54 @@ dependencies = [
 ]
 
 [[package]]
-name = "cubecl-hip"
-version = "0.10.0"
+name = "cubecl-environment"
+version = "0.11.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3c6b510a9348f06ecd56b32cf10eb6241639e927a87f5c09ec61cc7a7610d5a3"
+checksum = "61c7e23b2d90817886a9dc7c0a4527343def447ddbba1bae39f4f9f12c367385"
+dependencies = [
+ "async-channel",
+ "backtrace",
+ "bytemuck",
+ "bytes",
+ "cfg-if",
+ "cfg_aliases",
+ "ciborium",
+ "embassy-futures",
+ "embassy-time",
+ "etcetera",
+ "foldhash 0.2.0",
+ "futures-lite",
+ "hashbrown 0.17.1",
+ "log",
+ "oneshot",
+ "parking_lot",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rand 0.10.1",
+ "rusqlite",
+ "serde",
+ "serde_bytes",
+ "serde_json",
+ "spin",
+ "toml",
+ "tracing",
+ "wasm-bindgen-futures",
+ "web-time",
+]
+
+[[package]]
+name = "cubecl-hip"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7eaad6c93c5c48aa075457f03d8eeeca6f3ccb01f19a7da4940c664d48a14248"
 dependencies = [
  "bytemuck",
- "cubecl-common",
- "cubecl-core",
- "cubecl-cpp",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-core 0.11.0-pre.2",
+ "cubecl-cpp 0.11.0-pre.2",
+ "cubecl-environment",
  "cubecl-hip-sys",
- "cubecl-runtime",
+ "cubecl-runtime 0.11.0-pre.2",
  "derive-new",
  "half",
  "log",
@@ -1523,8 +1737,8 @@ version = "0.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d28a897a40c8ee6c57a8b8d76d8823ba2e97f4c2b83db9338f17a0752132d486"
 dependencies = [
- "cubecl-common",
- "cubecl-macros-internal",
+ "cubecl-common 0.10.0",
+ "cubecl-macros-internal 0.10.0",
  "derive-new",
  "derive_more",
  "enumset",
@@ -1536,7 +1750,33 @@ dependencies = [
  "num-traits",
  "portable-atomic",
  "serde",
- "variadics_please",
+ "variadics_please 1.1.0",
+]
+
+[[package]]
+name = "cubecl-ir"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "17ca2b43b9cfb6c4013592fa023a9f9158e344074d9dbba040b45461cc0842bc"
+dependencies = [
+ "bumpalo",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-macros-internal 0.11.0-pre.2",
+ "derive-new",
+ "derive_more",
+ "enumset",
+ "float-ord",
+ "fnv",
+ "foldhash 0.2.0",
+ "half",
+ "hashbrown 0.17.1",
+ "internment",
+ "itertools",
+ "num-traits",
+ "portable-atomic",
+ "serde",
+ "variadics_please 2.0.0",
 ]
 
 [[package]]
@@ -1545,7 +1785,7 @@ version = "0.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "77aa6df563e8e6a0926d2e9eeacb968737940a0e1ae9be819f6a65a341006e12"
 dependencies = [
- "cubecl-common",
+ "cubecl-common 0.10.0",
  "darling 0.23.0",
  "derive-new",
  "ident_case",
@@ -1556,6 +1796,24 @@ dependencies = [
  "syn 2.0.117",
 ]
 
+[[package]]
+name = "cubecl-macros"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f6b8fa7c11f7e311af88fc12883eedef8e0350541121fe8cc634db335c6e49ce"
+dependencies = [
+ "cubecl-common 0.11.0-pre.2",
+ "darling 0.23.0",
+ "derive-new",
+ "derive_more",
+ "ident_case",
+ "inflections",
+ "prettyplease",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.117",
+]
+
 [[package]]
 name = "cubecl-macros-internal"
 version = "0.10.0"
@@ -1568,15 +1826,27 @@ dependencies = [
  "syn 2.0.117",
 ]
 
+[[package]]
+name = "cubecl-macros-internal"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7082876a49f18dce60f20c34c42f6368d7753ea570238f2a4924fa0701fbfa82"
+dependencies = [
+ "darling 0.23.0",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.117",
+]
+
 [[package]]
 name = "cubecl-opt"
 version = "0.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0a599c6c3efefdcee8636994c90e58ef696c7efbed2d18b5b5ad8d5f29923abb"
 dependencies = [
- "cubecl-common",
- "cubecl-core",
- "cubecl-ir",
+ "cubecl-common 0.10.0",
+ "cubecl-core 0.10.0",
+ "cubecl-ir 0.10.0",
  "float-ord",
  "log",
  "num",
@@ -1586,6 +1856,26 @@ dependencies = [
  "type-map",
 ]
 
+[[package]]
+name = "cubecl-opt"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d0904a499f3ccbbfb24c2c8668d516f351c3e50c9178068be87d92bb3b8e1478"
+dependencies = [
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-core 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-ir 0.11.0-pre.2",
+ "float-ord",
+ "hashbrown 0.17.1",
+ "log",
+ "num",
+ "petgraph",
+ "smallvec",
+ "stable-vec",
+ "type-map",
+]
+
 [[package]]
 name = "cubecl-runtime"
 version = "0.10.0"
@@ -1597,9 +1887,9 @@ dependencies = [
  "bytemuck",
  "cfg-if",
  "cfg_aliases",
- "cubecl-common",
- "cubecl-ir",
- "cubecl-zspace",
+ "cubecl-common 0.10.0",
+ "cubecl-ir 0.10.0",
+ "cubecl-zspace 0.10.0",
  "derive-new",
  "derive_more",
  "dirs",
@@ -1612,26 +1902,55 @@ dependencies = [
  "spin",
  "thiserror 2.0.18",
  "toml",
- "tracing",
  "wasm-bindgen-futures",
  "web-time",
 ]
 
 [[package]]
-name = "cubecl-std"
-version = "0.10.0"
+name = "cubecl-runtime"
+version = "0.11.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b391b584a4897683dd9fc0767f96337ac712b733126ce0f68a9ee8a2df073d2d"
+checksum = "21231110bd2319a5a025d42749f03438a1a2caf18083f3a1b5bc7ab4554f32da"
 dependencies = [
- "cubecl-common",
- "cubecl-core",
- "cubecl-runtime",
+ "ahash",
+ "bytemuck",
+ "cfg-if",
+ "cfg_aliases",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-ir 0.11.0-pre.2",
+ "cubecl-zspace 0.11.0-pre.2",
+ "derive-new",
+ "derive_more",
+ "enumset",
+ "futures-util",
+ "hashbrown 0.17.1",
+ "itertools",
+ "log",
+ "md5",
+ "serde",
+ "serde_json",
+ "thiserror 2.0.18",
+ "toml",
+ "tracing",
+]
+
+[[package]]
+name = "cubecl-std"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bd029b15a150c3326b71a436f20b81973a2898ae9406a349c726918efdacf455"
+dependencies = [
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-core 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-runtime 0.11.0-pre.2",
  "half",
  "num-traits",
  "paste",
  "serde",
  "spin",
- "variadics_please",
+ "variadics_please 2.0.0",
 ]
 
 [[package]]
@@ -1644,19 +1963,44 @@ dependencies = [
  "bytemuck",
  "cfg-if",
  "cfg_aliases",
- "cubecl-common",
- "cubecl-core",
- "cubecl-ir",
- "cubecl-runtime",
+ "cubecl-common 0.10.0",
+ "cubecl-core 0.10.0",
+ "cubecl-ir 0.10.0",
+ "cubecl-runtime 0.10.0",
  "derive-new",
  "derive_more",
  "half",
  "hashbrown 0.16.1",
  "log",
  "sanitize-filename",
+ "wasm-bindgen-futures",
+ "wgpu 29.0.3",
+]
+
+[[package]]
+name = "cubecl-wgpu"
+version = "0.11.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aa87054aea892561fb0645491c5b156286d6d25c47e6f33fec7d2eeab002be5a"
+dependencies = [
+ "bytemuck",
+ "cfg-if",
+ "cfg_aliases",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-core 0.11.0-pre.2",
+ "cubecl-environment",
+ "cubecl-ir 0.11.0-pre.2",
+ "cubecl-runtime 0.11.0-pre.2",
+ "derive-new",
+ "derive_more",
+ "half",
+ "hashbrown 0.17.1",
+ "itertools",
+ "log",
+ "sanitize-filename",
  "tracing",
  "wasm-bindgen-futures",
- "wgpu",
+ "wgpu 30.0.0",
 ]
 
 [[package]]
@@ -1671,16 +2015,29 @@ dependencies = [
 ]
 
 [[package]]
-name = "cubek"
-version = "0.2.0"
+name = "cubecl-zspace"
+version = "0.11.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ba62fb5043d5eb723017415eec57c541f41901beea984500e2e16b57d75fa717"
+checksum = "f0513aa2d1997d1049305c5dae7f95bdcf808fc06ced02111075b46a67233950"
 dependencies = [
- "cubecl",
+ "derive-new",
+ "serde",
+ "smallvec",
+]
+
+[[package]]
+name = "cubek"
+version = "0.3.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f51262bc23900f4dc70c0a41f8d56910c3c6971e1dc242ec687cf920eda55c2"
+dependencies = [
+ "cubecl 0.11.0-pre.2",
  "cubek-attention",
  "cubek-convolution",
  "cubek-fft",
+ "cubek-interpolate",
  "cubek-matmul",
+ "cubek-pool",
  "cubek-quant",
  "cubek-random",
  "cubek-reduce",
@@ -1689,13 +2046,13 @@ dependencies = [
 
 [[package]]
 name = "cubek-attention"
-version = "0.2.0"
+version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "75ca5c363f05a9297b9d1c0e90959cc5f9fd90bc596e4e537f26f20602ec63a5"
+checksum = "739d905b83b5e42854cdfb8201389c688636043acc36c882b71a7ac10735e65b"
 dependencies = [
  "bytemuck",
- "cubecl",
- "cubecl-common",
+ "cubecl 0.11.0-pre.2",
+ "cubecl-common 0.11.0-pre.2",
  "cubek-matmul",
  "cubek-std",
  "half",
@@ -1704,13 +2061,13 @@ dependencies = [
 
 [[package]]
 name = "cubek-convolution"
-version = "0.2.0"
+version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "70f2755a8eb9ea5509c17edac78e4a9dc4be89c433d8e2ee709703a93b96f749"
+checksum = "b3c24e82dad0d0d0c6fcaf7d13376653529e91639ef401457e972f2e329fb018"
 dependencies = [
  "bytemuck",
- "cubecl",
- "cubecl-common",
+ "cubecl 0.11.0-pre.2",
+ "cubecl-common 0.11.0-pre.2",
  "cubek-matmul",
  "cubek-std",
  "derive-new",
@@ -1721,47 +2078,74 @@ dependencies = [
 
 [[package]]
 name = "cubek-fft"
-version = "0.2.0"
+version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8dd7c2972363332dc4609067a2ccc21856308d98089e676d3f0c52747ae61a5b"
+checksum = "ddf0f55951bfd0d0ceddb67fc14f34e7484ef6a1838b9522ad2bd36c5cba11c3"
 dependencies = [
- "cubecl",
+ "cubecl 0.11.0-pre.2",
+]
+
+[[package]]
+name = "cubek-interpolate"
+version = "0.3.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ce68e9c94bb9274827aca056c2589b24ce5acf2abba6aa8840bc27841eaad99"
+dependencies = [
+ "cubecl 0.11.0-pre.2",
+ "cubecl-common 0.11.0-pre.2",
+ "half",
+ "serde",
+ "thiserror 2.0.18",
 ]
 
 [[package]]
 name = "cubek-matmul"
-version = "0.2.0"
+version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "83a4cea5f0f439907dc953c7638a6204b3f055f1bcbd10db91dfc5faa030ac1c"
+checksum = "204917532e8f5bc4b440c640551d2c10d402bac7fd8062450d843752352518df"
 dependencies = [
  "bytemuck",
- "cubecl",
- "cubecl-common",
+ "cubecl 0.11.0-pre.2",
+ "cubecl-common 0.11.0-pre.2",
  "cubek-std",
+ "cubek-tile",
  "half",
  "serde",
 ]
 
 [[package]]
-name = "cubek-quant"
-version = "0.2.0"
+name = "cubek-pool"
+version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "802ff1b5f19597a7b7cb308c7c69e9eef28b78d57dab7794c816bb8f7d3d1b85"
+checksum = "f9e09ab3dd3a448b7672bc4ce4a880909896c733fb42e958c035d76dd6bbf291"
 dependencies = [
- "cubecl",
- "cubecl-common",
+ "cubecl 0.11.0-pre.2",
+ "cubecl-common 0.11.0-pre.2",
+ "thiserror 2.0.18",
+]
+
+[[package]]
+name = "cubek-quant"
+version = "0.3.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "76384f2666b4c552881e146af6fe543aecf422ecff4d9ffbefb75b91461cf0e5"
+dependencies = [
+ "cubecl 0.11.0-pre.2",
+ "cubecl-common 0.11.0-pre.2",
+ "cubek-tile",
  "half",
  "serde",
 ]
 
 [[package]]
 name = "cubek-random"
-version = "0.2.0"
+version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8dfecce959bdbeb3b355120e586d1ea9e45a0fc7f2ca354f4260a571952551be"
+checksum = "eb20538cb2a0d548a5571846021911d4ee660f75b1c10f6237c22cce1a2d4f7e"
 dependencies = [
- "cubecl",
- "cubecl-common",
+ "cubecl 0.11.0-pre.2",
+ "cubecl-common 0.11.0-pre.2",
+ "cubecl-environment",
  "half",
  "num-traits",
  "rand 0.10.1",
@@ -1770,11 +2154,11 @@ dependencies = [
 
 [[package]]
 name = "cubek-reduce"
-version = "0.2.0"
+version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "86758b84452b06e9121f25200da13b4909ecd6235cefb5f33f459bea24dac41a"
+checksum = "5f05c73af6d1b35f886661b44813985284d86f6ce6d3395c064f52c6e8d50f81"
 dependencies = [
- "cubecl",
+ "cubecl 0.11.0-pre.2",
  "cubek-std",
  "half",
  "num-traits",
@@ -1784,15 +2168,28 @@ dependencies = [
 
 [[package]]
 name = "cubek-std"
-version = "0.2.0"
+version = "0.3.0-pre.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a587d0a44b6f14f4c3711ac10717179b48adb59203b2b2bffe334876bf713187"
+checksum = "db70badd24b211310d9db3ccc280ccd148100f5f42a112c594fbc0e6b3c9a5dd"
 dependencies = [
- "cubecl",
- "cubecl-common",
+ "cubecl 0.11.0-pre.2",
+ "cubecl-common 0.11.0-pre.2",
  "half",
 ]
 
+[[package]]
+name = "cubek-tile"
+version = "0.3.0-pre.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b21695c6662edd669148874aea0eb8f0b011c7ef5bf15efa242b8bdae5d71555"
+dependencies = [
+ "bytemuck",
+ "cubecl 0.11.0-pre.2",
+ "cubecl-common 0.11.0-pre.2",
+ "half",
+ "serde",
+]
+
 [[package]]
 name = "cudarc"
 version = "0.19.7"
@@ -1905,15 +2302,6 @@ dependencies = [
  "syn 2.0.117",
 ]
 
-[[package]]
-name = "dary_heap"
-version = "0.3.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe"
-dependencies = [
- "serde",
-]
-
 [[package]]
 name = "data-encoding"
 version = "2.11.0"
@@ -1946,37 +2334,6 @@ dependencies = [
  "syn 2.0.117",
 ]
 
-[[package]]
-name = "derive_builder"
-version = "0.20.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947"
-dependencies = [
- "derive_builder_macro",
-]
-
-[[package]]
-name = "derive_builder_core"
-version = "0.20.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8"
-dependencies = [
- "darling 0.20.11",
- "proc-macro2",
- "quote",
- "syn 2.0.117",
-]
-
-[[package]]
-name = "derive_builder_macro"
-version = "0.20.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c"
-dependencies = [
- "derive_builder_core",
- "syn 2.0.117",
-]
-
 [[package]]
 name = "derive_more"
 version = "2.1.1"
@@ -2162,11 +2519,11 @@ name = "engine"
 version = "0.1.0"
 dependencies = [
  "burn",
- "burn-ndarray",
+ "burn-ndarray 0.21.0",
  "chess",
  "ctrlc",
  "rand 0.10.1",
- "rand_distr 0.6.0",
+ "rand_distr",
  "serde",
  "serde_json",
 ]
@@ -2248,10 +2605,14 @@ dependencies = [
 ]
 
 [[package]]
-name = "esaxx-rs"
-version = "0.1.10"
+name = "etcetera"
+version = "0.11.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6"
+checksum = "de48cc4d1c1d97a20fd819def54b890cadde72ed3ad0c614822a0a433361be96"
+dependencies = [
+ "cfg-if",
+ "windows-sys 0.61.2",
+]
 
 [[package]]
 name = "euclid"
@@ -2271,6 +2632,8 @@ dependencies = [
  "concurrent-queue",
  "parking",
  "pin-project-lite",
+ "portable-atomic",
+ "portable-atomic-util",
 ]
 
 [[package]]
@@ -2320,6 +2683,18 @@ dependencies = [
  "synstructure 0.12.6",
 ]
 
+[[package]]
+name = "fallible-iterator"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649"
+
+[[package]]
+name = "fallible-streaming-iterator"
+version = "0.1.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a"
+
 [[package]]
 name = "fancy-regex"
 version = "0.11.0"
@@ -2425,9 +2800,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c2d1f04709a8ac06e8e8042875a3c466cc4832d3c1a18dbcb9dba3c6e83046bc"
 dependencies = [
  "half",
- "num-traits",
- "rand 0.9.4",
- "rand_distr 0.5.1",
 ]
 
 [[package]]
@@ -2740,6 +3112,18 @@ version = "0.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280"
 
+[[package]]
+name = "gloo-timers"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbb143cf96099802033e0d4f4963b19fd2e0b728bcf076cd9cf7f6634f092994"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "js-sys",
+ "wasm-bindgen",
+]
+
 [[package]]
 name = "glow"
 version = "0.17.0"
@@ -2824,8 +3208,6 @@ dependencies = [
  "cfg-if",
  "crunchy",
  "num-traits",
- "rand 0.9.4",
- "rand_distr 0.5.1",
  "serde",
  "zerocopy",
 ]
@@ -2845,6 +3227,8 @@ version = "0.15.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1"
 dependencies = [
+ "allocator-api2",
+ "equivalent",
  "foldhash 0.1.5",
 ]
 
@@ -2866,6 +3250,22 @@ name = "hashbrown"
 version = "0.17.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
+dependencies = [
+ "allocator-api2",
+ "equivalent",
+ "foldhash 0.2.0",
+ "serde",
+ "serde_core",
+]
+
+[[package]]
+name = "hashlink"
+version = "0.12.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32069d97bb81e38fa67eab65e3393bf804bb85969f2bc06bf13f64aef5aba248"
+dependencies = [
+ "hashbrown 0.17.1",
+]
 
 [[package]]
 name = "heck"
@@ -2982,7 +3382,7 @@ version = "0.1.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0"
 dependencies = [
- "base64 0.22.1",
+ "base64",
  "bytes",
  "futures-channel",
  "futures-util",
@@ -3203,6 +3603,16 @@ dependencies = [
  "syn 2.0.117",
 ]
 
+[[package]]
+name = "internment"
+version = "0.8.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "636d4b0f6a39fd684effe2a73f5310df16a3fa7954c26d36833e98f44d1977a2"
+dependencies = [
+ "hashbrown 0.15.5",
+ "serde",
+]
+
 [[package]]
 name = "interpolate_name"
 version = "0.2.4"
@@ -3220,15 +3630,6 @@ version = "2.12.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
 
-[[package]]
-name = "itertools"
-version = "0.13.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186"
-dependencies = [
- "either",
-]
-
 [[package]]
 name = "itertools"
 version = "0.14.0"
@@ -3284,13 +3685,12 @@ dependencies = [
 
 [[package]]
 name = "js-sys"
-version = "0.3.98"
+version = "0.3.104"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "67df7112613f8bfd9150013a0314e196f4800d3201ae742489d999db2f979f08"
+checksum = "0e0c1080212aad755ea003d18543e8768dd432c48819efd73a7bf1e39b7a5a3a"
 dependencies = [
  "cfg-if",
  "futures-util",
- "once_cell",
  "wasm-bindgen",
 ]
 
@@ -3418,6 +3818,17 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "libsqlite3-sys"
+version = "0.38.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f1d20bef17f513b9b3004532233187769cd072d790971f4e4da0e346eb6401e8"
+dependencies = [
+ "cc",
+ "pkg-config",
+ "vcpkg",
+]
+
 [[package]]
 name = "line-clipping"
 version = "0.3.7"
@@ -3496,9 +3907,9 @@ dependencies = [
 
 [[package]]
 name = "macerator"
-version = "0.3.3"
+version = "0.3.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "508a2f720538bb7e3ea3cb6d098615b107cb82ae387d77d906276905e9ec894b"
+checksum = "2ddeaadb76e307b1d2b4836adee77910d405b9f6c1ca031bc81a7a398e2b22f2"
 dependencies = [
  "bytemuck",
  "cfg_aliases",
@@ -3512,9 +3923,9 @@ dependencies = [
 
 [[package]]
 name = "macerator-macros"
-version = "0.1.5"
+version = "0.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ed5ce85961d618ce9794bdf822bfe96fe9dd341aa5b033b454f7a8d96e79b9b1"
+checksum = "d2a397d020af346358830d6f28579f26283003ce87af3dd79173ac9ee2d3ac28"
 dependencies = [
  "darling 0.20.11",
  "proc-macro2",
@@ -3522,22 +3933,6 @@ dependencies = [
  "syn 2.0.117",
 ]
 
-[[package]]
-name = "macro_rules_attribute"
-version = "0.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520"
-dependencies = [
- "macro_rules_attribute-proc_macro",
- "paste",
-]
-
-[[package]]
-name = "macro_rules_attribute-proc_macro"
-version = "0.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30"
-
 [[package]]
 name = "matrixmultiply"
 version = "0.3.10"
@@ -3580,7 +3975,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3"
 dependencies = [
  "libc",
- "stable_deref_trait",
 ]
 
 [[package]]
@@ -3632,28 +4026,6 @@ version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4a0b3262dc837d2513fe2ef31ff8461352ef932dcca31ba0c0abe33547cf6b9b"
 
-[[package]]
-name = "monostate"
-version = "0.1.18"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67"
-dependencies = [
- "monostate-impl",
- "serde",
- "serde_core",
-]
-
-[[package]]
-name = "monostate-impl"
-version = "0.1.18"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.117",
-]
-
 [[package]]
 name = "moxcms"
 version = "0.8.1"
@@ -3664,6 +4036,12 @@ dependencies = [
  "pxfm",
 ]
 
+[[package]]
+name = "mutants"
+version = "0.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc0287524726960e07b119cebd01678f852f147742ae0d925e6a520dca956126"
+
 [[package]]
 name = "naga"
 version = "29.0.3"
@@ -3690,6 +4068,44 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "naga"
+version = "30.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23bf0a141a9ab6f07dbb492db53245e464bc9db42f407772d9ae03d83a2c1033"
+dependencies = [
+ "arrayvec 0.7.6",
+ "bit-set 0.10.0",
+ "bitflags 2.11.1",
+ "cfg-if",
+ "cfg_aliases",
+ "codespan-reporting",
+ "half",
+ "hashbrown 0.17.1",
+ "indexmap",
+ "libm",
+ "log",
+ "naga-types",
+ "num-traits",
+ "once_cell",
+ "rustc-hash 1.1.0",
+ "spirv",
+ "thiserror 2.0.18",
+ "unicode-ident",
+]
+
+[[package]]
+name = "naga-types"
+version = "30.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "590afbf58a6f4f62873cd5cff4468061844bafa1cdf399cc954537c22d768d49"
+dependencies = [
+ "hashbrown 0.17.1",
+ "indexmap",
+ "rustc-hash 1.1.0",
+ "thiserror 2.0.18",
+]
+
 [[package]]
 name = "nb"
 version = "0.1.3"
@@ -3965,6 +4381,18 @@ dependencies = [
  "objc2",
 ]
 
+[[package]]
+name = "objc2-core-graphics"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807"
+dependencies = [
+ "dispatch2",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-io-surface",
+]
+
 [[package]]
 name = "objc2-encode"
 version = "4.1.0"
@@ -3992,6 +4420,17 @@ dependencies = [
  "objc2-core-foundation",
 ]
 
+[[package]]
+name = "objc2-io-surface"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d"
+dependencies = [
+ "bitflags 2.11.1",
+ "objc2",
+ "objc2-core-foundation",
+]
+
 [[package]]
 name = "objc2-metal"
 version = "0.3.2"
@@ -4013,6 +4452,7 @@ dependencies = [
  "bitflags 2.11.1",
  "objc2",
  "objc2-core-foundation",
+ "objc2-core-graphics",
  "objc2-foundation",
  "objc2-metal",
 ]
@@ -4031,6 +4471,10 @@ name = "once_cell"
 version = "1.21.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+dependencies = [
+ "critical-section",
+ "portable-atomic",
+]
 
 [[package]]
 name = "oneshot"
@@ -4038,28 +4482,6 @@ version = "0.2.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cfe21416a02c693fb9f980befcb230ecc70b0b3d1cc4abf88b9675c4c1457f0c"
 
-[[package]]
-name = "onig"
-version = "6.5.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2"
-dependencies = [
- "bitflags 2.11.1",
- "libc",
- "once_cell",
- "onig_sys",
-]
-
-[[package]]
-name = "onig_sys"
-version = "69.9.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7"
-dependencies = [
- "cc",
- "pkg-config",
-]
-
 [[package]]
 name = "option-ext"
 version = "0.2.0"
@@ -4292,6 +4714,7 @@ version = "1.13.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
 dependencies = [
+ "critical-section",
  "serde",
 ]
 
@@ -4605,16 +5028,6 @@ version = "0.10.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69"
 
-[[package]]
-name = "rand_distr"
-version = "0.5.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463"
-dependencies = [
- "num-traits",
- "rand 0.9.4",
-]
-
 [[package]]
 name = "rand_distr"
 version = "0.6.0"
@@ -4673,7 +5086,7 @@ dependencies = [
  "compact_str",
  "hashbrown 0.16.1",
  "indoc",
- "itertools 0.14.0",
+ "itertools",
  "kasuari",
  "lru",
  "strum 0.27.2",
@@ -4725,7 +5138,7 @@ dependencies = [
  "hashbrown 0.16.1",
  "indoc",
  "instability",
- "itertools 0.14.0",
+ "itertools",
  "line-clipping",
  "ratatui-core",
  "strum 0.27.2",
@@ -4750,7 +5163,7 @@ dependencies = [
  "built",
  "cfg-if",
  "interpolate_name",
- "itertools 0.14.0",
+ "itertools",
  "libc",
  "libfuzzer-sys",
  "log",
@@ -4827,17 +5240,6 @@ dependencies = [
  "rayon-core",
 ]
 
-[[package]]
-name = "rayon-cond"
-version = "0.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f"
-dependencies = [
- "either",
- "itertools 0.14.0",
- "rayon",
-]
-
 [[package]]
 name = "rayon-core"
 version = "1.13.0"
@@ -4921,7 +5323,7 @@ version = "0.12.28"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
 dependencies = [
- "base64 0.22.1",
+ "base64",
  "bytes",
  "futures-channel",
  "futures-core",
@@ -4995,6 +5397,16 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "rsqlite-vfs"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c"
+dependencies = [
+ "hashbrown 0.16.1",
+ "thiserror 2.0.18",
+]
+
 [[package]]
 name = "rstest"
 version = "0.26.1"
@@ -5024,6 +5436,21 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "rusqlite"
+version = "0.40.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23f2a97da3e3873c73cb2a2e71b35c40ff95e0b1eefa8d72d8499a6928c3b5b3"
+dependencies = [
+ "bitflags 2.11.1",
+ "fallible-iterator",
+ "fallible-streaming-iterator",
+ "hashlink",
+ "libsqlite3-sys",
+ "smallvec",
+ "sqlite-wasm-rs",
+]
+
 [[package]]
 name = "rustc-demangle"
 version = "0.1.27"
@@ -5396,15 +5823,15 @@ dependencies = [
 ]
 
 [[package]]
-name = "spm_precompiled"
-version = "0.1.4"
+name = "sqlite-wasm-rs"
+version = "0.5.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326"
+checksum = "dc3efc0da82635d7e1ced0053bbbfa8c7ab9645d0bf36ceb4f7127bb85315d75"
 dependencies = [
- "base64 0.13.1",
- "nom 7.1.3",
- "serde",
- "unicode-segmentation",
+ "cc",
+ "js-sys",
+ "rsqlite-vfs",
+ "wasm-bindgen",
 ]
 
 [[package]]
@@ -5507,6 +5934,17 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "syn"
+version = "3.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
 [[package]]
 name = "sync_wrapper"
 version = "1.0.2"
@@ -5592,7 +6030,7 @@ dependencies = [
  "safetensors 0.3.3",
  "thiserror 1.0.69",
  "torch-sys",
- "zip 0.6.6",
+ "zip",
 ]
 
 [[package]]
@@ -5645,7 +6083,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4676b37242ccbd1aabf56edb093a4827dc49086c0ffd764a5705899e0f35f8f7"
 dependencies = [
  "anyhow",
- "base64 0.22.1",
+ "base64",
  "bitflags 2.11.1",
  "fancy-regex",
  "filedescriptor",
@@ -5836,39 +6274,6 @@ version = "0.1.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
 
-[[package]]
-name = "tokenizers"
-version = "0.22.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223"
-dependencies = [
- "ahash",
- "aho-corasick",
- "compact_str",
- "dary_heap",
- "derive_builder",
- "esaxx-rs",
- "getrandom 0.3.4",
- "itertools 0.14.0",
- "log",
- "macro_rules_attribute",
- "monostate",
- "onig",
- "paste",
- "rand 0.9.4",
- "rayon",
- "rayon-cond",
- "regex",
- "regex-syntax",
- "serde",
- "serde_json",
- "spm_precompiled",
- "thiserror 2.0.18",
- "unicode-normalization-alignments",
- "unicode-segmentation",
- "unicode_categories",
-]
-
 [[package]]
 name = "tokio"
 version = "1.52.3"
@@ -5879,10 +6284,23 @@ dependencies = [
  "libc",
  "mio",
  "pin-project-lite",
+ "signal-hook-registry",
  "socket2",
+ "tokio-macros",
  "windows-sys 0.61.2",
 ]
 
+[[package]]
+name = "tokio-macros"
+version = "2.7.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 3.0.4",
+]
+
 [[package]]
 name = "tokio-rustls"
 version = "0.26.4"
@@ -5969,7 +6387,7 @@ dependencies = [
  "serde",
  "serde_json",
  "ureq",
- "zip 0.6.6",
+ "zip",
 ]
 
 [[package]]
@@ -6019,9 +6437,9 @@ checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3"
 
 [[package]]
 name = "tracel-llvm"
-version = "20.1.4-7"
+version = "22.1.4-5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "982535db9eb1a30ac0f2c50239a0eec3e5cf50993a88e92b04747bd2f4d365b2"
+checksum = "1fdf8a19f43dfaa86f20c7e7f22376b1bb47f4846cd89f4e0bb9f1a4fa50a9e9"
 dependencies = [
  "tracel-mlir-rs",
  "tracel-mlir-sys",
@@ -6029,9 +6447,9 @@ dependencies = [
 
 [[package]]
 name = "tracel-llvm-bundler"
-version = "20.1.4-7"
+version = "22.1.4-5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3c75b8e477cb8d49d907afab029ca74d48459f5b88c27bdb4c6cd6acb5e61977"
+checksum = "51a3f7a4f173e6ccec871f31fba2d4cb71c5c41ab0089c0bf9559f4780dfaf07"
 dependencies = [
  "anyhow",
  "bytes",
@@ -6049,9 +6467,9 @@ dependencies = [
 
 [[package]]
 name = "tracel-mlir-rs"
-version = "20.1.4-7"
+version = "22.1.4-5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "77a478a35efd68d0ba73f747adfb7923b121c64e7f5be9cd8364ca1dcb772d5c"
+checksum = "b0b1cf4683d1bf518040a8e13ebb2c1b18a8c2bb1db257f47603edbd99b094f8"
 dependencies = [
  "tracel-mlir-rs-macros",
  "tracel-mlir-sys",
@@ -6059,9 +6477,9 @@ dependencies = [
 
 [[package]]
 name = "tracel-mlir-rs-macros"
-version = "20.1.4-7"
+version = "22.1.4-5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7a94f36868c3b10b1825945223d99d106c73f4d249f063caa4651deeb9379344"
+checksum = "04d4cea4e89fcaddbcc7fdeba0720b62d8bda63e23326ff28463d4b864ef9945"
 dependencies = [
  "comrak",
  "convert_case 0.8.0",
@@ -6076,21 +6494,19 @@ dependencies = [
 
 [[package]]
 name = "tracel-mlir-sys"
-version = "20.1.4-7"
+version = "22.1.4-5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "02f26d31af0c225a6d2e3d65d012fd6de848c9fc776897b152ee83b7d1bd15c4"
+checksum = "fb813a9f259bb2112a5a1186fa366e6075e6dc1c041b5ae38a116535b9daca64"
 dependencies = [
  "tracel-llvm-bundler",
 ]
 
 [[package]]
 name = "tracel-tblgen-rs"
-version = "20.1.4-7"
+version = "22.1.4-5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "00d2581070380418ccc33b500f3739e4d4869421fdb477fcea51ff97c6253a52"
+checksum = "d93699345d51e8bcb0d968165bf024ea28af87fee776a60d55875bdb1b898f2d"
 dependencies = [
- "bindgen",
- "cc",
  "paste",
  "thiserror 2.0.18",
  "tracel-llvm-bundler",
@@ -6196,12 +6612,6 @@ version = "2.0.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6af6ae20167a9ece4bcb41af5b80f8a1f1df981f6391189ce00fd257af04126a"
 
-[[package]]
-name = "typed-path"
-version = "0.12.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8e28f89b80c87b8fb0cf04ab448d5dd0dd0ade2f8891bae878de66a75a28600e"
-
 [[package]]
 name = "typenum"
 version = "1.20.0"
@@ -6237,15 +6647,6 @@ dependencies = [
  "tinyvec",
 ]
 
-[[package]]
-name = "unicode-normalization-alignments"
-version = "0.1.12"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de"
-dependencies = [
- "smallvec",
-]
-
 [[package]]
 name = "unicode-segmentation"
 version = "1.13.2"
@@ -6258,7 +6659,7 @@ version = "2.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "16b380a1238663e5f8a691f9039c73e1cdae598a30e9855f541d29b08b53e9a5"
 dependencies = [
- "itertools 0.14.0",
+ "itertools",
  "unicode-segmentation",
  "unicode-width",
 ]
@@ -6287,25 +6688,30 @@ version = "0.2.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
 
+[[package]]
+name = "unsynn"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "501a7adf1a4bd9951501e5c66621e972ef8874d787628b7f90e64f936ef7ec0a"
+dependencies = [
+ "mutants",
+ "proc-macro2",
+ "rustc-hash 2.1.2",
+]
+
 [[package]]
 name = "untrusted"
 version = "0.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
 
-[[package]]
-name = "unty"
-version = "0.0.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6d49784317cd0d1ee7ec5c716dd598ec5b4483ea832a2dced265471cc0f690ae"
-
 [[package]]
 name = "ureq"
 version = "2.12.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d"
 dependencies = [
- "base64 0.22.1",
+ "base64",
  "flate2",
  "log",
  "once_cell",
@@ -6381,6 +6787,22 @@ dependencies = [
  "syn 2.0.117",
 ]
 
+[[package]]
+name = "variadics_please"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b53d0f7f2d0182759a549f1f313ce16b07d8f9ba93098157b5fa10ee3e61117f"
+dependencies = [
+ "quote",
+ "unsynn",
+]
+
+[[package]]
+name = "vcpkg"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
+
 [[package]]
 name = "version_check"
 version = "0.9.5"
@@ -6447,9 +6869,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen"
-version = "0.2.121"
+version = "0.2.127"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "49ace1d07c165b0864824eee619580c4689389afa9dc9ed3a4c75040d82e6790"
+checksum = "1b70935747edd64d89de3efa29d73789b806c15798f8e7dca4d8ac356b50ce70"
 dependencies = [
  "cfg-if",
  "once_cell",
@@ -6460,9 +6882,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-futures"
-version = "0.4.71"
+version = "0.4.77"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "96492d0d3ffba25305a7dc88720d250b1401d7edca02cc3bcd50633b424673b8"
+checksum = "6b7777d5cc23d0e91404e53ce2d5e8ec7acae3026b16233dba62cd3246457950"
 dependencies = [
  "js-sys",
  "wasm-bindgen",
@@ -6470,9 +6892,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro"
-version = "0.2.121"
+version = "0.2.127"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8e68e6f4afd367a562002c05637acb8578ff2dea1943df76afb9e83d177c8578"
+checksum = "77775f8f3f7217702089053b94958f8f54061a3f663417df76e19cbdcca29bc1"
 dependencies = [
  "quote",
  "wasm-bindgen-macro-support",
@@ -6480,9 +6902,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-macro-support"
-version = "0.2.121"
+version = "0.2.127"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d95a9ec35c64b2a7cb35d3fead40c4238d0940c86d107136999567a4703259f2"
+checksum = "e11d33f857dc2fb11b8bc75aee111aa9cbeb12cd9f25efd3d4c2a3dd4e235284"
 dependencies = [
  "bumpalo",
  "proc-macro2",
@@ -6493,9 +6915,9 @@ dependencies = [
 
 [[package]]
 name = "wasm-bindgen-shared"
-version = "0.2.121"
+version = "0.2.127"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c4e0100b01e9f0d03189a92b96772a1fb998639d981193d7dbab487302513441"
+checksum = "7ef64dbcc55df09c7e5a46182d181c2cfa3e925f3da937ea764728b4bbb9dcbf"
 dependencies = [
  "unicode-ident",
 ]
@@ -6552,9 +6974,9 @@ version = "0.1.0"
 
 [[package]]
 name = "web-sys"
-version = "0.3.98"
+version = "0.3.104"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4b572dff8bcf38bad0fa19729c89bb5748b2b9b1d8be70cf90df697e3a8f32aa"
+checksum = "c435338968042f4f59a557f690a253676d47ce13ceb55d70100e7facf6620a30"
 dependencies = [
  "js-sys",
  "wasm-bindgen",
@@ -6681,7 +7103,7 @@ dependencies = [
  "hashbrown 0.16.1",
  "js-sys",
  "log",
- "naga",
+ "naga 29.0.3",
  "parking_lot",
  "portable-atomic",
  "profiling",
@@ -6691,9 +7113,39 @@ dependencies = [
  "wasm-bindgen",
  "wasm-bindgen-futures",
  "web-sys",
- "wgpu-core",
- "wgpu-hal",
- "wgpu-types",
+ "wgpu-core 29.0.3",
+ "wgpu-hal 29.0.3",
+ "wgpu-types 29.0.3",
+]
+
+[[package]]
+name = "wgpu"
+version = "30.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d8f4bd44d92da5270f03409dba9f952dab24f128e05d6a554926101d1bf9114"
+dependencies = [
+ "arrayvec 0.7.6",
+ "bitflags 2.11.1",
+ "bytemuck",
+ "cfg-if",
+ "cfg_aliases",
+ "document-features",
+ "hashbrown 0.17.1",
+ "js-sys",
+ "log",
+ "naga 30.0.0",
+ "parking_lot",
+ "portable-atomic",
+ "profiling",
+ "raw-window-handle",
+ "smallvec",
+ "static_assertions",
+ "wasm-bindgen",
+ "wasm-bindgen-futures",
+ "web-sys",
+ "wgpu-core 30.0.0",
+ "wgpu-hal 30.0.0",
+ "wgpu-types 30.0.1",
 ]
 
 [[package]]
@@ -6712,7 +7164,7 @@ dependencies = [
  "hashbrown 0.16.1",
  "indexmap",
  "log",
- "naga",
+ "naga 29.0.3",
  "once_cell",
  "parking_lot",
  "portable-atomic",
@@ -6721,12 +7173,46 @@ dependencies = [
  "rustc-hash 1.1.0",
  "smallvec",
  "thiserror 2.0.18",
- "wgpu-core-deps-apple",
- "wgpu-core-deps-emscripten",
- "wgpu-core-deps-windows-linux-android",
- "wgpu-hal",
- "wgpu-naga-bridge",
- "wgpu-types",
+ "wgpu-core-deps-apple 29.0.3",
+ "wgpu-core-deps-emscripten 29.0.3",
+ "wgpu-core-deps-windows-linux-android 29.0.3",
+ "wgpu-hal 29.0.3",
+ "wgpu-naga-bridge 29.0.3",
+ "wgpu-types 29.0.3",
+]
+
+[[package]]
+name = "wgpu-core"
+version = "30.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08763620e76fc980bca7bf84de82568614487a53172dd968d89187282eb87fa2"
+dependencies = [
+ "arrayvec 0.7.6",
+ "bit-set 0.10.0",
+ "bit-vec 0.9.1",
+ "bitflags 2.11.1",
+ "bytemuck",
+ "cfg_aliases",
+ "document-features",
+ "hashbrown 0.17.1",
+ "indexmap",
+ "log",
+ "naga 30.0.0",
+ "naga-types",
+ "once_cell",
+ "parking_lot",
+ "portable-atomic",
+ "profiling",
+ "raw-window-handle",
+ "rustc-hash 1.1.0",
+ "smallvec",
+ "thiserror 2.0.18",
+ "wgpu-core-deps-apple 30.0.1",
+ "wgpu-core-deps-emscripten 30.0.1",
+ "wgpu-core-deps-windows-linux-android 30.0.1",
+ "wgpu-hal 30.0.0",
+ "wgpu-naga-bridge 30.0.1",
+ "wgpu-types 30.0.1",
 ]
 
 [[package]]
@@ -6735,7 +7221,16 @@ version = "29.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "62e51b5447e144b3dbba4feb01f80f4fa21696fa0cd99afb2c3df1affd6fdb28"
 dependencies = [
- "wgpu-hal",
+ "wgpu-hal 29.0.3",
+]
+
+[[package]]
+name = "wgpu-core-deps-apple"
+version = "30.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "061f3d319a40d39d00b1ecc2c33b89fe21d4e6fe01859df3500a3a8ecccd6b68"
+dependencies = [
+ "wgpu-hal 30.0.0",
 ]
 
 [[package]]
@@ -6744,7 +7239,16 @@ version = "29.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3487cd6293a963bc5c0c0396f6a2192043c50003c07f4efdccbad3d90ec9d819"
 dependencies = [
- "wgpu-hal",
+ "wgpu-hal 29.0.3",
+]
+
+[[package]]
+name = "wgpu-core-deps-emscripten"
+version = "30.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d98b86cf4abf524a902dd35f18ca6a3f08fc2ae9847c8f10b48e30491b1f0b86"
+dependencies = [
+ "wgpu-hal 30.0.0",
 ]
 
 [[package]]
@@ -6753,7 +7257,16 @@ version = "29.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1bfb01076d0aa08b0ba9bd741e178b5cc440f5abe99d9581323a4c8b5d1a1916"
 dependencies = [
- "wgpu-hal",
+ "wgpu-hal 29.0.3",
+]
+
+[[package]]
+name = "wgpu-core-deps-windows-linux-android"
+version = "30.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7586165fd5f6d881cb9ce4bb71f40d6caab2c0f1837e3fc1d9788a197fb6004f"
+dependencies = [
+ "wgpu-hal 30.0.0",
 ]
 
 [[package]]
@@ -6781,7 +7294,7 @@ dependencies = [
  "libc",
  "libloading 0.8.9",
  "log",
- "naga",
+ "naga 29.0.3",
  "ndk-sys",
  "objc2",
  "objc2-core-foundation",
@@ -6803,8 +7316,64 @@ dependencies = [
  "wasm-bindgen",
  "wayland-sys",
  "web-sys",
- "wgpu-naga-bridge",
- "wgpu-types",
+ "wgpu-naga-bridge 29.0.3",
+ "wgpu-types 29.0.3",
+ "windows",
+ "windows-core",
+ "windows-result",
+]
+
+[[package]]
+name = "wgpu-hal"
+version = "30.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf765132d8d5f50e192e7880464890c13f4e7457aafe8e5466e8174586e9f101"
+dependencies = [
+ "android_system_properties",
+ "arrayvec 0.7.6",
+ "ash",
+ "bit-set 0.10.0",
+ "bitflags 2.11.1",
+ "block2",
+ "bytemuck",
+ "cfg-if",
+ "cfg_aliases",
+ "glow",
+ "glutin_wgl_sys",
+ "gpu-allocator",
+ "hashbrown 0.17.1",
+ "js-sys",
+ "khronos-egl",
+ "libc",
+ "libloading 0.8.9",
+ "log",
+ "naga 30.0.0",
+ "naga-types",
+ "ndk-sys",
+ "objc2",
+ "objc2-core-foundation",
+ "objc2-core-graphics",
+ "objc2-foundation",
+ "objc2-metal",
+ "objc2-quartz-core",
+ "once_cell",
+ "ordered-float 5.3.0",
+ "parking_lot",
+ "portable-atomic",
+ "portable-atomic-util",
+ "profiling",
+ "range-alloc",
+ "raw-window-handle",
+ "raw-window-metal",
+ "renderdoc-sys",
+ "smallvec",
+ "static_assertions",
+ "thiserror 2.0.18",
+ "wasm-bindgen",
+ "wayland-sys",
+ "web-sys",
+ "wgpu-naga-bridge 30.0.1",
+ "wgpu-types 30.0.1",
  "windows",
  "windows-core",
  "windows-result",
@@ -6816,8 +7385,18 @@ version = "29.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "59c654c483f058800972c3645e95388a7eca31bf9fe1933bc20e036588a0be02"
 dependencies = [
- "naga",
- "wgpu-types",
+ "naga 29.0.3",
+ "wgpu-types 29.0.3",
+]
+
+[[package]]
+name = "wgpu-naga-bridge"
+version = "30.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d2f62e73117bb7a62bfd9c5a5841438a823f6566c6442a808ee269d2d055c081"
+dependencies = [
+ "naga 30.0.0",
+ "wgpu-types 30.0.1",
 ]
 
 [[package]]
@@ -6834,6 +7413,22 @@ dependencies = [
  "web-sys",
 ]
 
+[[package]]
+name = "wgpu-types"
+version = "30.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "99dad6f1fbdbbdb4c278a6508b059d44688f5cebddf78d005a46a31340269286"
+dependencies = [
+ "bitflags 2.11.1",
+ "bytemuck",
+ "js-sys",
+ "log",
+ "naga-types",
+ "raw-window-handle",
+ "static_assertions",
+ "web-sys",
+]
+
 [[package]]
 name = "winapi"
 version = "0.3.9"
@@ -7391,18 +7986,6 @@ dependencies = [
  "zstd",
 ]
 
-[[package]]
-name = "zip"
-version = "7.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c42e33efc22a0650c311c2ef19115ce232583abbe80850bc8b66509ebef02de0"
-dependencies = [
- "crc32fast",
- "indexmap",
- "memchr",
- "typed-path",
-]
-
 [[package]]
 name = "zmij"
 version = "1.0.21"
diff --git a/engine/Cargo.toml b/engine/Cargo.toml
index a2752c9..56d8f70 100644
--- a/engine/Cargo.toml
+++ b/engine/Cargo.toml
@@ -5,7 +5,7 @@ edition = "2024"
 
 [dependencies]
 chess = "3.2.0"
-burn = { version = "0.21.0", features = ["std", "tui", "train", "cuda", "fusion", "wgpu", "ndarray"], default-features = false }
+burn = { version = "0.22.0-pre.2", features = ["std", "tui", "train", "cuda", "fusion", "wgpu", "ndarray"], default-features = false }
 burn-ndarray = "0.21.0"
 rand_distr = "0.6.0"
 rand = "0.10.1"
diff --git a/engine/src/lib.rs b/engine/src/lib.rs
index 8f87ede..02effbd 100644
--- a/engine/src/lib.rs
+++ b/engine/src/lib.rs
@@ -1,7 +1,6 @@
 pub mod mcts;
 mod net;
-mod train;
-pub mod training;
+pub mod train;
 
 use chess::Game;
 
diff --git a/engine/src/main.rs b/engine/src/main.rs
index 999aa86..50a8a21 100644
--- a/engine/src/main.rs
+++ b/engine/src/main.rs
@@ -1,9 +1,9 @@
 #![recursion_limit = "256"]
 
-use burn::backend::{Autodiff, Cuda};
 use burn::optim::AdamConfig;
+use burn::tensor::{Device, DeviceIndex};
 use engine::mcts::MctsConfig;
-use engine::training::train::{train, TrainingConfig};
+use engine::train::{train, TrainingConfig};
 // fn main() {
 //     type MyBackend = Wgpu<f32, i32>;
 //
@@ -14,13 +14,7 @@ use engine::training::train::{train, TrainingConfig};
 // }
 
 fn main() {
-    // type MyBackend = Wgpu<f32, i32>;
-    type MyBackend = Cuda<f32, i32>;
-    // type MyBackend = NdArray<f32, i32>;
-    type MyAutodiffBackend = Autodiff<MyBackend>;
-    // let device = burn::backend::wgpu::WgpuDevice::default();
-    // let device = burn::backend::ndarray::NdArrayDevice::default();
-    let device = burn::backend::cuda::CudaDevice::default();
+    let device = Device::cuda(DeviceIndex::default());
 
     let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
 
@@ -43,7 +37,7 @@ fn main() {
         seed: None,
     };
 
-    train::<MyAutodiffBackend>(training_config, device);
+    train(training_config, &device);
 
     // let mut board_state: Option<BoardState> = Some(BoardState::new(
     //     Board::default(),
diff --git a/engine/src/mcts.rs b/engine/src/mcts.rs
index 288818e..7d49c82 100644
--- a/engine/src/mcts.rs
+++ b/engine/src/mcts.rs
@@ -2,7 +2,7 @@
 use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
 use crate::net::encoding::{encode_board_state_perspective, encode_move};
 use crate::net::model::ChessModel;
-use burn::prelude::Backend;
+use burn::prelude::Device;
 use burn::Tensor;
 use chess::BoardStatus::{Checkmate, Stalemate};
 use chess::Color::{Black, White};
@@ -10,7 +10,6 @@ use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
 use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
 use rand::SeedableRng;
 use std::collections::HashMap;
-use std::marker::PhantomData;
 
 pub struct Node {
     pub prior: f32,
@@ -107,17 +106,62 @@ impl Default for MctsConfig {
     }
 }
 
-pub struct Mcts<B: Backend> {
+pub struct Mcts {
     pub config: MctsConfig,
-    pub _marker: PhantomData<B>, // if B not otherwise stored
 }
 
-impl<B: Backend> Mcts<B> {
+fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
+    match status {
+        BoardStateStatus::Ongoing => None,
+        BoardStateStatus::Stalemate
+        | BoardStateStatus::Threefold
+        | BoardStateStatus::FiftyMove => Some(0.0),
+        BoardStateStatus::WhiteWinner => Some(if stm == White { 1.0 } else { -1.0 }),
+        BoardStateStatus::BlackWinner => Some(if stm == Black { 1.0 } else { -1.0 }),
+    }
+}
+
+fn softmax(logits: &[f32]) -> Vec<f32> {
+    if logits.is_empty() {
+        return Vec::new();
+    }
+    let mut max_logit = f32::NEG_INFINITY;
+    for &v in logits {
+        if v > max_logit {
+            max_logit = v;
+        }
+    }
+    if !max_logit.is_finite() {
+        max_logit = 0.0;
+    }
+    let mut exps_sum = 0.0f32;
+    let mut probs = vec![0.0f32; logits.len()];
+    for (i, &v) in logits.iter().enumerate() {
+        let x = (v - max_logit).exp();
+        if x.is_finite() {
+            probs[i] = x;
+            exps_sum += x;
+        }
+    }
+    if exps_sum > 0.0 && exps_sum.is_finite() {
+        for p in &mut probs {
+            *p /= exps_sum;
+        }
+    } else {
+        let uniform = 1.0 / logits.len() as f32;
+        for p in &mut probs {
+            *p = uniform;
+        }
+    }
+    probs
+}
+
+impl Mcts {
     pub fn search(
         &mut self,
         board_state: &BoardState,
-        model: &ChessModel<B>,
-        device: &B::Device,
+        model: &ChessModel,
+        device: &Device,
     ) -> MctsResults {
         let mut nodes = Vec::<Node>::new();
 
@@ -130,133 +174,138 @@ impl<B: Backend> Mcts<B> {
         // Apply Dirichlet noise to root children
         self.add_dirichlet_noise(root, &mut nodes);
 
-        // We'll batch leaf evaluations to reduce per-leaf model calls and device-host syncs.
         let mut sims_done: usize = 0;
         let num_sims = self.config.num_simulations;
 
         while sims_done < num_sims {
-            // Collect a batch of leaf nodes (and their selection paths)
             let mut leaf_nodes: Vec<usize> = Vec::new();
             let mut leaf_paths: Vec<Vec<usize>> = Vec::new();
-            let mut leaf_states: Vec<Tensor<B, 4>> = Vec::new();
+            let mut leaf_states: Vec<Tensor<4>> = Vec::new();
 
-            while leaf_nodes.len() < std::cmp::min(self.config.batch_max, num_sims - sims_done) {
+            let target_batch = std::cmp::min(self.config.batch_max, num_sims - sims_done);
+
+            while leaf_nodes.len() < target_batch && sims_done + leaf_nodes.len() < num_sims {
                 let mut path = vec![root];
                 let mut current = root;
 
+                // Apply virtual loss during selection path
                 while !nodes[current].children.is_empty() {
+                    nodes[current].visit_count += 1;
+                    nodes[current].value_sum -= 1.0;
+
                     current = nodes[current].select_child(&nodes, &self.config.c_puct);
                     path.push(current);
                 }
 
-                // Record leaf node and its path
-                leaf_nodes.push(current);
-                leaf_paths.push(path.clone());
+                // Check terminal status of leaf
+                let stm = nodes[current].board_state.board.side_to_move();
+                if let Some(term_val) = terminal_value(nodes[current].board_state.status, stm) {
+                    // Revert virtual loss for path except leaf itself
+                    for &idx in &path[..path.len() - 1] {
+                        nodes[idx].visit_count -= 1;
+                        nodes[idx].value_sum += 1.0;
+                    }
+                    self.backpropagate(&mut nodes, &path, term_val, stm);
+                    sims_done += 1;
+                    continue;
+                }
 
-                // Prepare state tensor for this leaf
-                let state: Tensor<B, 4> =
+                // Prevent expanding the exact same node twice in a single batch
+                if leaf_nodes.contains(&current) {
+                    // Revert virtual loss for this duplicate path
+                    for &idx in &path[..path.len() - 1] {
+                        nodes[idx].visit_count -= 1;
+                        nodes[idx].value_sum += 1.0;
+                    }
+                    break;
+                }
+
+                leaf_nodes.push(current);
+                leaf_paths.push(path);
+
+                let state: Tensor<4> =
                     encode_board_state_perspective(&nodes[current].board_state, device)
                         .reshape([1, 18, 8, 8]);
                 leaf_states.push(state);
-
-                sims_done += 1;
             }
 
             if leaf_nodes.is_empty() {
-                break;
+                if sims_done >= num_sims {
+                    break;
+                }
+                continue;
             }
 
-            // Batch evaluate the collected leaf states
+            // Batch evaluate non-terminal leaf states
             let batch = Tensor::cat(leaf_states, 0);
             let (policy_batch, value_batch) = model.forward(batch);
 
-            // Move tensors to host once per batch
             let policy_data = policy_batch.into_data().to_vec::<f32>().unwrap();
             let value_data = value_batch.into_data().to_vec::<f32>().unwrap();
 
             let num_moves = policy_data.len() / leaf_nodes.len();
 
-            // Process each evaluated leaf: expand and backpropagate
             for (i, &node_idx) in leaf_nodes.iter().enumerate() {
                 let path = &leaf_paths[i];
 
-                // slice for this sample's logits
+                // Revert virtual loss along this path before applying true expansion & backpropagation
+                for &idx in &path[..path.len() - 1] {
+                    nodes[idx].visit_count -= 1;
+                    nodes[idx].value_sum += 1.0;
+                }
+
                 let start = i * num_moves;
                 let end = start + num_moves;
                 let logits = &policy_data[start..end];
+                let probs = softmax(logits);
 
-                // Convert logits to probabilities with a numerically-stable softmax on host
-                let mut max_logit = f32::NEG_INFINITY;
-                for &v in logits.iter() {
-                    if v > max_logit {
-                        max_logit = v;
-                    }
-                }
-                if logits.is_empty() {
-                    println!("logits empty")
-                }
-                if !max_logit.is_finite() {
-                    println!("max logits not finite")
-                }
-
-                let mut exps_sum = 0.0f32;
-                let mut probs = vec![0.0f32; num_moves];
-
-                for (j, &v) in logits.iter().enumerate() {
-                    let x = (v - max_logit).exp();
-
-                    if !x.is_finite() {
-                        probs[j] = 0.0;
-                    } else {
-                        probs[j] = x;
-                        exps_sum += x;
-                    }
-                }
-
-                if exps_sum > 0.0 && exps_sum.is_finite() {
-                    for p in &mut probs {
-                        *p /= exps_sum;
-                    }
-                } else {
-                    // fallback: uniform or zero
-                    let uniform = 1.0 / num_moves as f32;
-                    for p in &mut probs {
-                        *p = uniform;
-                    }
-                }
-
-                // Expand: add legal moves as children with prior from probs
+                let stm = nodes[node_idx].board_state.board.side_to_move();
                 let legal_moves: Vec<ChessMove> =
                     MoveGen::new_legal(&nodes[node_idx].board_state.board).collect();
-                for mv in legal_moves {
-                    let stm = nodes[node_idx].board_state.board.side_to_move();
-                    let idx = encode_move(mv, stm).expect("Invalid move");
-                    let prior = probs[idx];
+
+                let mut move_priors = Vec::with_capacity(legal_moves.len());
+                let mut prior_sum = 0.0f32;
+
+                for &mv in &legal_moves {
+                    if let Some(enc_idx) = encode_move(mv, stm) {
+                        let p = probs[enc_idx];
+                        move_priors.push((mv, p));
+                        prior_sum += p;
+                    }
+                }
+
+                for (mv, p) in move_priors {
+                    let normalized_prior = if prior_sum > 0.0 {
+                        p / prior_sum
+                    } else {
+                        1.0 / legal_moves.len() as f32
+                    };
 
                     let mut new_board = nodes[node_idx].board_state.clone();
                     new_board.apply_move(mv);
 
                     let child_idx = nodes.len();
-                    nodes.push(Node::new(prior, new_board, Some(mv)));
+                    nodes.push(Node::new(normalized_prior, new_board, Some(mv)));
                     nodes[node_idx].children.push(child_idx);
                 }
 
-                // Backpropagate the value for this leaf
                 let value = value_data[i];
-                let color = nodes[node_idx].board_state.board.side_to_move();
-                self.backpropagate(&mut nodes, path, value, color);
+                self.backpropagate(&mut nodes, path, value, stm);
+                sims_done += 1;
             }
         }
 
-        // Build compact move distribution: encoded move index -> probability (visits / num_simulations)
+        // Build compact move distribution
         let mut move_dist: Vec<(usize, f32)> = Vec::with_capacity(nodes[root].children.len());
         let stm = board_state.board.side_to_move();
-        let denom = self.config.num_simulations as f32;
+        let total_root_visits: u32 = nodes[root].children.iter().map(|&c| nodes[c].visit_count).sum();
+        let denom = if total_root_visits > 0 { total_root_visits as f32 } else { 1.0 };
         for child_idx in nodes[root].children.iter() {
             let mv = nodes[*child_idx].last_move.expect("move didnt exist");
-            let enc = encode_move(mv, stm).expect("Invalid move");
-            let prob = nodes[*child_idx].visit_count as f32 / denom;
-            move_dist.push((enc, prob));
+            if let Some(enc) = encode_move(mv, stm) {
+                let prob = nodes[*child_idx].visit_count as f32 / denom;
+                move_dist.push((enc, prob));
+            }
         }
 
         MctsResults::new(board_state.clone(), move_dist, nodes[root].value())
@@ -266,54 +315,53 @@ impl<B: Backend> Mcts<B> {
         &mut self,
         node_idx: usize,
         arena: &mut Vec<Node>,
-        model: &ChessModel<B>,
-        device: &B::Device,
+        model: &ChessModel,
+        device: &Device,
     ) -> f32 {
-        if arena[node_idx].board_state.status == BoardStateStatus::Stalemate
-            || arena[node_idx].board_state.status == BoardStateStatus::Threefold
-            || arena[node_idx].board_state.status == BoardStateStatus::FiftyMove
-        {
-            0.0
-        } else if arena[node_idx].board_state.status == BoardStateStatus::WhiteWinner {
-            if arena[node_idx].board_state.board.side_to_move() == Black {
-                1.0
-            } else {
-                -1.0
-            }
-        } else if arena[node_idx].board_state.status == BoardStateStatus::BlackWinner {
-            if arena[node_idx].board_state.board.side_to_move() == White {
-                1.0
-            } else {
-                -1.0
-            }
+        let status = arena[node_idx].board_state.status;
+        let stm = arena[node_idx].board_state.board.side_to_move();
+
+        if let Some(term_val) = terminal_value(status, stm) {
+            term_val
         } else {
-            let state: Tensor<B, 4> =
+            let state: Tensor<4> =
                 encode_board_state_perspective(&arena[node_idx].board_state, device)
                     .reshape([1, 18, 8, 8]);
-            // let start = Instant::now();
             let (policy_head, value_head) = model.forward(state);
-            // println!("time: {:?}", start.elapsed());
 
             let legal_moves: Vec<ChessMove> =
                 MoveGen::new_legal(&arena[node_idx].board_state.board).collect();
 
-            let policy = policy_head.into_data().to_vec::<f32>().unwrap();
+            let logits = policy_head.into_data().to_vec::<f32>().unwrap();
+            let probs = softmax(&logits);
 
-            for mv in legal_moves {
-                let stm = arena[node_idx].board_state.board.side_to_move();
-                let idx = encode_move(mv, stm).expect("Invalid move");
-                let prior = policy[idx];
+            let mut move_priors = Vec::with_capacity(legal_moves.len());
+            let mut prior_sum = 0.0f32;
+
+            for &mv in &legal_moves {
+                if let Some(enc_idx) = encode_move(mv, stm) {
+                    let p = probs[enc_idx];
+                    move_priors.push((mv, p));
+                    prior_sum += p;
+                }
+            }
+
+            for (mv, p) in move_priors {
+                let normalized_prior = if prior_sum > 0.0 {
+                    p / prior_sum
+                } else {
+                    1.0 / legal_moves.len() as f32
+                };
 
                 let mut new_board = arena[node_idx].board_state.clone();
                 new_board.apply_move(mv);
 
                 let child_idx = arena.len();
-
-                arena.push(Node::new(prior, new_board, Some(mv)));
+                arena.push(Node::new(normalized_prior, new_board, Some(mv)));
                 arena[node_idx].children.push(child_idx);
             }
 
-            value_head.into_data().to_vec().unwrap()[0]
+            value_head.into_data().to_vec::<f32>().unwrap()[0]
         }
     }
 
@@ -356,7 +404,6 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
 
     let gamma = Gamma::new(alpha as f64, 1.0).unwrap();
 
-    // Use a single SmallRng seeded from system time (avoid depending on thread_rng helper)
     let now = std::time::SystemTime::now()
         .duration_since(std::time::UNIX_EPOCH)
         .unwrap();
@@ -367,8 +414,15 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
 
     let sum: f32 = samples.iter().sum();
 
-    for x in &mut samples {
-        *x /= sum;
+    if sum > 0.0 && sum.is_finite() {
+        for x in &mut samples {
+            *x /= sum;
+        }
+    } else if size > 0 {
+        let uniform = 1.0 / size as f32;
+        for x in &mut samples {
+            *x = uniform;
+        }
     }
 
     samples
diff --git a/engine/src/net/encoding.rs b/engine/src/net/encoding.rs
index b1526ba..7593c8c 100644
--- a/engine/src/net/encoding.rs
+++ b/engine/src/net/encoding.rs
@@ -52,10 +52,7 @@ const UNDERPROMO_DIRS: [(i8, i8); 3] = [
 
 const UNDERPROMO_PIECES: [Piece; 3] = [Piece::Knight, Piece::Bishop, Piece::Rook];
 
-pub fn encode_board_state_perspective<B: Backend>(
-    state: &BoardState,
-    device: &B::Device,
-) -> Tensor<B, 3> {
+pub fn encode_board_state_perspective(state: &BoardState, device: &Device) -> Tensor<3> {
     let mut planes = vec![0.0f32; 18 * 64];
 
     let board = &state.board;
@@ -142,7 +139,7 @@ pub fn encode_board_state_perspective<B: Backend>(
         planes[idx(17, rank, file)] = 1.0;
     }
 
-    Tensor::<B, 3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
+    Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
 }
 
 fn fill_plane(buffer: &mut [f32], plane: usize) {
@@ -201,7 +198,7 @@ pub fn encode_move(mv: ChessMove, stm: Color) -> Option<usize> {
     };
 
     let from_sq = (fr as usize) * 8 + (ff as usize);
-    Some(from_sq * PLANES_PER_SQUARE + plane)
+    Some(plane * 64 + from_sq)
 }
 
 fn encode_non_underpromo(dx: i8, dy: i8) -> Option<usize> {
@@ -228,13 +225,13 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
         return None;
     }
 
-    let from_idx = idx / PLANES_PER_SQUARE;
-    let plane = idx % PLANES_PER_SQUARE;
+    let plane = idx / 64;
+    let from_idx = idx % 64;
 
     let ff = (from_idx % 8) as i8;
     let fr = (from_idx / 8) as i8;
 
-    let (dx, dy, promo) = if plane < 56 {
+    let (dx, dy, mut promo) = if plane < 56 {
         let dir = plane / 7;
         let dist = (plane % 7 + 1) as i8;
 
@@ -260,6 +257,11 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
         return None;
     }
 
+    // Default pawn moves to final rank (perspective fr == 6 -> tr == 7) to Queen promotion if not underpromotion
+    if fr == 6 && tr == 7 && dy == 1 && dx.abs() <= 1 && promo.is_none() {
+        promo = Some(Piece::Queen);
+    }
+
     let from = deorient_square(ff, fr, stm);
     let to = deorient_square(tf, tr, stm);
 
@@ -311,4 +313,26 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn queen_promotion_roundtrip() {
+        use chess::{File, Rank, Square};
+        let from = Square::make_square(Rank::Seventh, File::E);
+        let to = Square::make_square(Rank::Eighth, File::E);
+        let mv = ChessMove::new(from, to, Some(Piece::Queen));
+        let enc = encode_move(mv, Color::White).unwrap();
+        let dec = decode_move(enc, Color::White).unwrap();
+        assert_eq!(mv, dec);
+    }
+
+    #[test]
+    fn underpromotion_roundtrip() {
+        use chess::{File, Rank, Square};
+        let from = Square::make_square(Rank::Seventh, File::E);
+        let to = Square::make_square(Rank::Eighth, File::E);
+        let mv = ChessMove::new(from, to, Some(Piece::Knight));
+        let enc = encode_move(mv, Color::White).unwrap();
+        let dec = decode_move(enc, Color::White).unwrap();
+        assert_eq!(mv, dec);
+    }
 }
diff --git a/engine/src/net/model.rs b/engine/src/net/model.rs
index 12cd024..8a6666c 100644
--- a/engine/src/net/model.rs
+++ b/engine/src/net/model.rs
@@ -5,14 +5,12 @@ use burn::nn::conv::Conv2dConfig;
 use burn::nn::loss::{MseLoss, Reduction};
 use burn::nn::{BatchNorm, BatchNormConfig, LinearConfig, PaddingConfig2d};
 use burn::tensor::activation::log_softmax;
-use burn::tensor::backend::AutodiffBackend;
 use burn::tensor::Transaction;
 use burn::train::{InferenceStep, ItemLazy, TrainOutput, TrainStep};
 use burn::{
     nn::{conv::Conv2d, Linear, Relu},
     prelude::*,
 };
-use burn_ndarray::NdArray;
 /*
 Input planes:
 1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
@@ -39,11 +37,11 @@ pub struct ModelMetadata {
 }
 
 #[derive(Module, Debug)]
-pub struct ResidualBlock<B: Backend> {
-    conv1: Conv2d<B>,
-    bn1: BatchNorm<B>,
-    conv2: Conv2d<B>,
-    bn2: BatchNorm<B>,
+pub struct ResidualBlock {
+    conv1: Conv2d,
+    bn1: BatchNorm,
+    conv2: Conv2d,
+    bn2: BatchNorm,
     activation: Relu,
 }
 
@@ -52,8 +50,8 @@ pub struct ResidualBlockConfig {
     channels: usize,
 }
 
-impl<B: Backend> ResidualBlock<B> {
-    pub fn new(channels: usize, device: &B::Device) -> Self {
+impl ResidualBlock {
+    pub fn new(channels: usize, device: &Device) -> Self {
         Self {
             conv1: Conv2dConfig::new([channels, channels], [3, 3])
                 .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
@@ -66,7 +64,7 @@ impl<B: Backend> ResidualBlock<B> {
             activation: Relu::new(),
         }
     }
-    pub fn forward(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
+    pub fn forward(&self, x: Tensor<4>) -> Tensor<4> {
         let residual = x.clone();
 
         let x = self.conv1.forward(x);
@@ -81,24 +79,24 @@ impl<B: Backend> ResidualBlock<B> {
 }
 
 #[derive(Module, Debug)]
-pub struct ChessModel<B: Backend> {
+pub struct ChessModel {
     // Initial convolution
-    conv: Conv2d<B>,
-    bn: BatchNorm<B>,
+    conv: Conv2d,
+    bn: BatchNorm,
 
     // Residual tower
-    residual_blocks: Vec<ResidualBlock<B>>,
+    residual_blocks: Vec<ResidualBlock>,
 
     // Policy head
-    policy_conv: Conv2d<B>,
-    policy_bn: BatchNorm<B>,
-    policy_fc: Linear<B>,
+    policy_conv: Conv2d,
+    policy_bn: BatchNorm,
+    policy_conv2: Conv2d,
 
     // Value head
-    value_conv: Conv2d<B>,
-    value_bn: BatchNorm<B>,
-    value_fc1: Linear<B>,
-    value_fc2: Linear<B>,
+    value_conv: Conv2d,
+    value_bn: BatchNorm,
+    value_fc1: Linear,
+    value_fc2: Linear,
 
     activation: Relu,
 }
@@ -110,11 +108,7 @@ pub struct ChessModelConfig {
 }
 
 impl ChessModelConfig {
-    pub fn init<B: Backend>(
-        num_blocks: usize,
-        channels: usize,
-        device: &B::Device,
-    ) -> ChessModel<B> {
+    pub fn init(num_blocks: usize, channels: usize, device: &Device) -> ChessModel {
         ChessModel {
             conv: Conv2dConfig::new([18, channels], [3, 3]) // 18 plane input
                 .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
@@ -126,14 +120,14 @@ impl ChessModelConfig {
                 .collect(),
 
             // Policy head
-            policy_conv: Conv2dConfig::new([channels, 2], [1, 1]).init(device),
-            policy_bn: BatchNormConfig::new(2).init(device),
-            policy_fc: LinearConfig::new(2 * 8 * 8, 8 * 8 * 73).init(device), // 4672 typical chess move space
+            policy_conv: Conv2dConfig::new([channels, 32], [1, 1]).init(device),
+            policy_bn: BatchNormConfig::new(32).init(device),
+            policy_conv2: Conv2dConfig::new([32, 73], [1, 1]).init(device),
 
             // Value head
             value_conv: Conv2dConfig::new([channels, 1], [1, 1]).init(device),
             value_bn: BatchNormConfig::new(1).init(device),
-            value_fc1: LinearConfig::new(1 * 8 * 8, 256).init(device),
+            value_fc1: LinearConfig::new(8 * 8, 256).init(device),
             value_fc2: LinearConfig::new(256, 1).init(device),
 
             activation: Relu::new(),
@@ -141,8 +135,8 @@ impl ChessModelConfig {
     }
 }
 
-impl<B: Backend> ChessModel<B> {
-    pub fn forward(&self, x: Tensor<B, 4>) -> (Tensor<B, 2>, Tensor<B, 2>) {
+impl ChessModel {
+    pub fn forward(&self, x: Tensor<4>) -> (Tensor<2>, Tensor<2>) {
         let mut x = self.conv.forward(x);
         x = self.bn.forward(x);
         x = self.activation.forward(x);
@@ -157,8 +151,8 @@ impl<B: Backend> ChessModel<B> {
         let mut p = self.policy_conv.forward(x.clone());
         p = self.policy_bn.forward(p);
         p = self.activation.forward(p);
-        let mut p = p.reshape([batch_size, 2 * 8 * 8]);
-        p = self.policy_fc.forward(p);
+        p = self.policy_conv2.forward(p);
+        let p = p.reshape([batch_size, 73 * 8 * 8]);
 
         // -------- Value Head --------
         let mut v = self.value_conv.forward(x);
@@ -173,10 +167,10 @@ impl<B: Backend> ChessModel<B> {
 
     pub fn forward_chess(
         &self,
-        boards: Tensor<B, 4>,         // e.g. [batch, channels, 8, 8]
-        policy_targets: Tensor<B, 2>, // move distribution (make sure is normalized)
-        value_targets: Tensor<B, 2>,  // scalar evaluation
-    ) -> ChessOutput<B> {
+        boards: Tensor<4>,         // e.g. [batch, channels, 8, 8]
+        policy_targets: Tensor<2>, // move distribution (make sure is normalized)
+        value_targets: Tensor<2>,  // scalar evaluation
+    ) -> ChessOutput {
         let (policy_logits, value) = self.forward(boards);
 
         let log_probs = log_softmax(policy_logits.clone(), 1);
@@ -190,39 +184,44 @@ impl<B: Backend> ChessModel<B> {
         let value_loss =
             MseLoss::new().forward(value.clone(), value_targets.clone(), Reduction::Mean);
 
-        let total_loss = policy_loss + 0.5 * value_loss;
+        let total_loss = policy_loss.clone() + 0.5 * value_loss.clone();
 
         ChessOutput {
             policy_logits,
             value,
             policy_targets,
             value_targets,
+            policy_loss,
+            value_loss,
             loss: total_loss,
         }
     }
 }
 
-pub struct ChessOutput<B: Backend> {
-    pub policy_logits: Tensor<B, 2>, // [sample, num_moves (4672)]
-    pub value: Tensor<B, 2>,         // [sample, 1 value]
-    pub policy_targets: Tensor<B, 2>,
-    pub value_targets: Tensor<B, 2>, // kept 2d for consistency?
-    pub loss: Tensor<B, 1>,          // [sample]
+pub struct ChessOutput {
+    pub policy_logits: Tensor<2>, // [sample, num_moves (4672)]
+    pub value: Tensor<2>,         // [sample, 1 value]
+    pub policy_targets: Tensor<2>,
+    pub value_targets: Tensor<2>, // kept 2d for consistency?
+    pub policy_loss: Tensor<1>,
+    pub value_loss: Tensor<1>,
+    pub loss: Tensor<1>,          // [sample]
 }
 
-impl<B: Backend> ItemLazy for ChessOutput<B> {
-    type ItemSync = ChessOutput<NdArray>;
-
-    fn sync(self) -> Self::ItemSync {
-        let [policy_logits, value, policy_targets, value_targets, loss] = Transaction::default()
-            .register(self.policy_logits)
-            .register(self.value)
-            .register(self.policy_targets)
-            .register(self.value_targets)
-            .register(self.loss)
-            .execute()
-            .try_into()
-            .expect("Correct amount of tensor data");
+impl ItemLazy for ChessOutput {
+    fn sync(self) -> Self {
+        let [policy_logits, value, policy_targets, value_targets, policy_loss, value_loss, loss] =
+            Transaction::default()
+                .register(self.policy_logits)
+                .register(self.value)
+                .register(self.policy_targets)
+                .register(self.value_targets)
+                .register(self.policy_loss)
+                .register(self.value_loss)
+                .register(self.loss)
+                .execute()
+                .try_into()
+                .expect("Correct amount of tensor data");
 
         let device = &Default::default();
 
@@ -231,27 +230,29 @@ impl<B: Backend> ItemLazy for ChessOutput<B> {
             value: Tensor::from_data(value, device),
             policy_targets: Tensor::from_data(policy_targets, device),
             value_targets: Tensor::from_data(value_targets, device),
+            policy_loss: Tensor::from_data(policy_loss, device),
+            value_loss: Tensor::from_data(value_loss, device),
             loss: Tensor::from_data(loss, device),
         }
     }
 }
 
-impl<B: AutodiffBackend> TrainStep for ChessModel<B> {
-    type Input = ChessBatch<B>;
-    type Output = ChessOutput<B>;
+impl TrainStep for ChessModel {
+    type Input = ChessBatch;
+    type Output = ChessOutput;
 
-    fn step(&self, batch: ChessBatch<B>) -> TrainOutput<ChessOutput<B>> {
+    fn step(&self, batch: ChessBatch) -> TrainOutput<ChessOutput> {
         let item = self.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
 
         TrainOutput::new(self, item.loss.backward(), item)
     }
 }
 
-impl<B: Backend> InferenceStep for ChessModel<B> {
-    type Input = ChessBatch<B>;
-    type Output = ChessOutput<B>;
+impl InferenceStep for ChessModel {
+    type Input = ChessBatch;
+    type Output = ChessOutput;
 
-    fn step(&self, batch: ChessBatch<B>) -> ChessOutput<B> {
+    fn step(&self, batch: ChessBatch) -> ChessOutput {
         self.forward_chess(batch.states, batch.policy_targets, batch.value_targets)
     }
 }
@@ -287,14 +288,14 @@ impl TrainingSample {
 pub struct ChessBatcher {}
 
 #[derive(Clone, Debug)]
-pub struct ChessBatch<B: Backend> {
-    pub states: Tensor<B, 4>,
-    pub policy_targets: Tensor<B, 2>,
-    pub value_targets: Tensor<B, 2>,
+pub struct ChessBatch {
+    pub states: Tensor<4>,
+    pub policy_targets: Tensor<2>,
+    pub value_targets: Tensor<2>,
 }
 
-impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
-    fn batch(&self, items: Vec<TrainingSample>, device: &B::Device) -> ChessBatch<B> {
+impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
+    fn batch(&self, items: Vec<TrainingSample>, device: &Device) -> ChessBatch {
         let state_tensors = items
             .iter()
             .map(|item| {
@@ -318,18 +319,15 @@ impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
                         *p /= sum;
                     }
                 }
-                Tensor::<B, 2>::from_floats(TensorData::new(policy, [1, 4672]), device).unsqueeze()
+                Tensor::<2>::from_floats(TensorData::new(policy, [1, 4672]), device)
             })
             .collect::<Vec<_>>();
 
         let value_target_tensors = items
             .iter()
             .map(|item| {
-                Tensor::<B, 2>::from_floats(
-                    TensorData::new(vec![item.value_target], [1, 1]),
-                    device,
-                )
-                .reshape([1, 1])
+                Tensor::<2>::from_floats(TensorData::new(vec![item.value_target], [1, 1]), device)
+                    .reshape([1, 1])
             })
             .collect::<Vec<_>>();
 
diff --git a/engine/src/rl/classification.rs b/engine/src/rl/classification.rs
deleted file mode 100644
index 60fde81..0000000
--- a/engine/src/rl/classification.rs
+++ /dev/null
@@ -1,59 +0,0 @@
-/// Multi-label classification output adapted for multiple metrics.
-///
-/// Supported metrics:
-/// - HammingScore
-/// - Precision (via ConfusionStatsInput)
-/// - Recall (via ConfusionStatsInput)
-/// - FBetaScore (via ConfusionStatsInput)
-/// - Loss
-#[derive(new)]
-pub struct MultiLabelSoftClassificationOutput<B: Backend> {
-    /// The loss.
-    pub loss: Tensor<B, 1>,
-
-    /// The label logits or probabilities. Shape: \[batch_size, num_classes\].
-    pub output: Tensor<B, 2>,
-
-    /// The ground truth labels (target values). Shape: \[batch_size, num_classes\].
-    pub targets: Tensor<B, 2>,
-}
-
-impl<B: Backend> ItemLazy for MultiLabelSoftClassificationOutput<B> {
-    type ItemSync = MultiLabelSoftClassificationOutput<NdArray>;
-
-    fn sync(self) -> Self::ItemSync {
-        let [output, loss, targets] = Transaction::default()
-            .register(self.output)
-            .register(self.loss)
-            .register(self.targets)
-            .execute()
-            .try_into()
-            .expect("Correct amount of tensor data");
-
-        let device = &Default::default();
-
-        MultiLabelSoftClassificationOutput {
-            output: Tensor::from_data(output, device),
-            loss: Tensor::from_data(loss, device),
-            targets: Tensor::from_data(targets, device),
-        }
-    }
-}
-
-impl<B: Backend> Adaptor<HammingScoreInput<B>> for MultiLabelSoftClassificationOutput<B> {
-    fn adapt(&self) -> HammingScoreInput<B> {
-        HammingScoreInput::new(self.output.clone(), self.targets.clone())
-    }
-}
-
-impl<B: Backend> Adaptor<LossInput<B>> for MultiLabelSoftClassificationOutput<B> {
-    fn adapt(&self) -> LossInput<B> {
-        LossInput::new(self.loss.clone())
-    }
-}
-
-impl<B: Backend> Adaptor<ConfusionStatsInput<B>> for MultiLabelSoftClassificationOutput<B> {
-    fn adapt(&self) -> ConfusionStatsInput<B> {
-        ConfusionStatsInput::new(self.output.clone(), self.targets.clone().bool())
-    }
-}
diff --git a/engine/src/train.rs b/engine/src/train.rs
index 0bf9df9..625ec5a 100644
--- a/engine/src/train.rs
+++ b/engine/src/train.rs
@@ -1,134 +1,314 @@
-// use crate::model::{ChessBatcher, ChessModel, ChessModelConfig};
-// use burn::config::Config;
-// use burn::data::dataloader::batcher::Batcher;
-// use burn::optim::{Adam, AdamConfig, Optimizer, SimpleOptimizer};
-// use burn::optim::adaptor::OptimizerAdaptor;
-// use burn::prelude::Backend;
-// use burn::tensor::backend::AutodiffBackend;
-// use burn::train::TrainStep;
-// use crate::mcts::MctsResults;
-//
-// #[derive(Config, Debug)]
-// pub struct ChessTrainerConfig {
-//     pub model: ChessModelConfig,
-//     pub optimizer: AdamConfig,
-//     #[config(default = 10)]
-//     pub num_epochs: usize,
-//     #[config(default = 64)]
-//     pub batch_size: usize,
-//     #[config(default = 4)]
-//     pub num_workers: usize,
-//     #[config(default = 42)]
-//     pub seed: u64,
-//     #[config(default = 1.0e-4)]
-//     pub learning_rate: f64,
-// }
-//
-// impl ChessTrainerConfig {
-//     pub fn init<B: Backend>(
-//         model_config: ChessModelConfig,
-//         optimizer: AdamConfig,
-//         num_epochs: usize,
-//         batch_size: usize,
-//         num_workers: usize,
-//         seed: u64,
-//         learning_rate: f64,
-//     ) -> ChessTrainer<B> {
-//         ChessTrainer {
-//             model: model_config::init(),
-//             optimizer: optimizer.init(),
-//             num_epochs,
-//
-//
-//         }
-//     }
-// }
-//
-// pub struct ChessTrainer<B: AutodiffBackend> {
-//     pub model: ChessModel<B>,
-//     pub optimizer: Adam,
-//     learning_rate: f64,
-//     pub batcher: ChessBatcher,
-//     pub device: B::Device,
-// }
-//
-// impl<B: AutodiffBackend> ChessTrainer<B> {
-//     pub fn new(model: ChessModel<B>, device: B::Device) -> Self {
-//         let optimizer = AdamConfig::new().init();
-//
-//         Self {
-//             model,
-//             optimizer,
-//             batcher: ChessBatcher::default(),
-//             device,
-//         }
-//     }
-//
-//     pub fn train_step(&mut self, batch_data: Vec<MctsResults>) -> f32 {
-//         // 1. Convert to tensors
-//         let batch = self.batcher.batch(batch_data, &self.device);
-//
-//         // 2. Forward + backward (your TrainStep impl)
-//         let output = self.model.step(batch);
-//
-//         // 3. Update weights
-//         self.optimizer.step(, &mut self.model, output.grads);
-//
-//         // 4. Return loss (for logging)
-//         let loss_tensor = output.item.loss.clone().into_data();
-//         let loss = loss_tensor.to_vec::<f32>().unwrap()[0];
-//
-//         loss
-//     }
-// }
+use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
+use crate::net::encoding::decode_move;
+use crate::net::model::{
+    ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
+};
+use burn::data::dataloader::batcher::Batcher;
+use burn::module::{AutodiffModule, Module};
+use burn::optim::{AdamConfig, GradientsParams, Optimizer};
+// use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
+use burn::prelude::Device;
+use burn::store::ModuleRecord;
+use chess::{ChessMove, Color};
+use rand::rngs::SmallRng;
+use rand::seq::SliceRandom;
+use rand::{RngExt, SeedableRng};
+use std::collections::VecDeque;
+use std::fs::File;
+use std::io::BufReader;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::time::Instant;
+use std::time::{SystemTime, UNIX_EPOCH};
 
-// fn create_artifact_dir(artifact_dir: &str) {
-//     // Remove existing artifacts before to get an accurate learner summary
-//     std::fs::remove_dir_all(artifact_dir).ok();
-//     std::fs::create_dir_all(artifact_dir).ok();
-// }
-//
-// pub fn train<B: AutodiffBackend>(
-//     artifact_dir: &str,
-//     config: ChessTrainingConfig,
-//     device: B::Device,
-// ) {
-//     create_artifact_dir(artifact_dir);
-//     config
-//         .save(format!("{artifact_dir}/config.json"))
-//         .expect("Config should be saved successfully");
-//
-//     B::seed(&device, config.seed);
-//
-//     let batcher = ChessBatcher::default();
+pub struct TrainingConfig {
+    pub max_time_s: Option<u64>,
+    pub num_iters: Option<usize>,
+    pub max_depth: u16, // unused
+    pub model_name: String,
+    pub load_model: bool,
+    pub hidden_channels: usize,
+    pub num_blocks: usize,
+    pub batch_size: usize,      // num positions sampled to a batch (2048)
+    pub num_episodes: usize,    // num games generated per iteration (5000)
+    pub buffer_max_size: usize, // max number of samples in the buffer
+    pub mcts_config: MctsConfig,
+    pub optimizer: AdamConfig,
+    pub lr: f64,
+    pub seed: Option<u64>,
+}
 
-// let dataloader_train = DataLoaderBuilder::new(batcher.clone())
-//     .batch_size(config.batch_size)
-//     .shuffle(config.seed)
-//     .num_workers(config.num_workers)
-//     .build(MnistDataset::train());
-//
-// let dataloader_test = DataLoaderBuilder::new(batcher)
-//     .batch_size(config.batch_size)
-//     .shuffle(config.seed)
-//     .num_workers(config.num_workers)
-//     .build(MnistDataset::test());
-//
-// let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test)
-//     .metrics((AccuracyMetric::new(), LossMetric::new()))
-//     .with_file_checkpointer(CompactRecorder::new())
-//     .num_epochs(config.num_epochs)
-//     .summary();
-//
-// let model = config.model.init::<B>(&device);
-// let result = training.launch(Learner::new(
-//     model,
-//     config.optimizer.init(),
-//     config.learning_rate,
-// ));
-//
-// result
-//     .model
-//     .save_file(format!("{artifact_dir}/model"), &CompactRecorder::new())
-//     .expect("Trained model should be saved successfully");
-// }
+pub fn train(training_config: TrainingConfig, device: &Device) {
+    let model_path = format!("artifacts/{}", training_config.model_name.as_str());
+    let metadata_path = format!("{}.json", model_path);
+    println!("Creating model...");
+    let mut model: ChessModel = ChessModelConfig::init(
+        training_config.num_blocks,
+        training_config.hidden_channels,
+        device,
+    );
+    if training_config.load_model {
+        println!("Loading model {}...", model_path);
+        // Load model in full precision from MessagePack file
+        let record = ModuleRecord::load(&model_path).expect("Failed to load model.");
+        model = model.load_record(record)
+    }
+
+    let train = Arc::new(AtomicBool::new(true));
+    let train_signal = Arc::clone(&train);
+
+    let mut iter: usize = 0;
+
+    if training_config.load_model {
+        let file = File::open(&metadata_path).unwrap();
+
+        // 2. Use a buffered reader for efficiency
+        let reader = BufReader::new(file);
+
+        // 3. Deserialize JSON directly from the file reader
+        let metadata: ModelMetadata = serde_json::from_reader(reader).unwrap();
+        iter = metadata.iterations;
+        println!("Loaded model had {} iters", iter);
+    }
+
+    let start_time = Instant::now();
+
+    ctrlc::set_handler(move || {
+        train_signal.store(false, Ordering::Relaxed);
+        println!("Finishing batch before exiting...");
+    })
+    .expect("Error setting Ctrl-C handler");
+
+    let mut replay_buffer: VecDeque<TrainingSample> = VecDeque::new();
+
+    let mut mcts: Mcts = Mcts {
+        config: training_config.mcts_config,
+    };
+
+    // Create RNG once and reuse it for sampling and shuffling
+    // Seed from system time (platform default entropy may be unavailable in some contexts)
+    let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
+    let seed = training_config.seed.unwrap_or(now.as_nanos() as u64);
+    println!("using seed: {}", seed);
+    let mut rng = SmallRng::seed_from_u64(seed);
+
+    // Initialize optimizer once so state (moments) persist across steps
+    let mut optim = training_config.optimizer.init();
+
+    println!("Starting training...");
+    while train.load(Ordering::Relaxed) {
+        println!("Iteration: {}", iter);
+        let infer_model = model.valid();
+        // Gen samples
+        println!("Generating {} games...", training_config.num_episodes);
+        for episode in 0..training_config.num_episodes {
+            println!("Episode: {}", episode);
+            let mut board_state = BoardState::default();
+            let mut episode_buffer: Vec<MctsResults> = vec![];
+
+            let mut game_hist = String::new();
+
+            let mut ply_count: usize = 0;
+
+            while board_state.status == BoardStateStatus::Ongoing {
+                // let before = Instant::now();
+                let results = mcts.search(&board_state, &infer_model, &device);
+                // println!("{}", before.elapsed().as_millis());
+                episode_buffer.push(results);
+
+                let temp = if ply_count < 30 {
+                    1.0
+                } else {
+                    0.0
+                };
+
+                let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
+
+                let stm = board_state.board.side_to_move();
+                let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
+                // println!("playing move: {}", mv);
+                if episode == 0 {
+                    game_hist.push_str(mv.to_string().as_str());
+                    game_hist.push(' ');
+                }
+
+                board_state.apply_move(mv);
+                ply_count += 1;
+                if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing {
+                    board_state.status = BoardStateStatus::FiftyMove;
+                }
+            }
+            if episode == 0 {
+                println!("Game history of first game of iteration: {}", game_hist);
+            }
+
+            for result in episode_buffer.iter() {
+                let stm = result.board_state.board.side_to_move();
+                let outcome = match board_state.status {
+                    BoardStateStatus::WhiteWinner => {
+                        if stm == Color::White {
+                            1.0
+                        } else {
+                            -1.0
+                        }
+                    }
+                    BoardStateStatus::BlackWinner => {
+                        if stm == Color::Black {
+                            1.0
+                        } else {
+                            -1.0
+                        }
+                    }
+                    _ => 0.0,
+                };
+
+                replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
+                    result.clone(),
+                    outcome,
+                ));
+
+                if replay_buffer.len() > training_config.buffer_max_size {
+                    replay_buffer.pop_front();
+                }
+            }
+        }
+        // train
+
+        if replay_buffer.is_empty() {
+            println!("Replay buffer is empty, skipping training step.");
+            continue;
+        }
+
+        let num_samples = std::cmp::min(training_config.batch_size, replay_buffer.len());
+        println!(
+            "Finished self-play! Training on {} samples (buffer size: {})...",
+            num_samples,
+            replay_buffer.len()
+        );
+        let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
+        indices.shuffle(&mut rng);
+
+        let samples = indices
+            .into_iter()
+            .take(num_samples)
+            .map(|i| replay_buffer[i].clone())
+            .collect::<Vec<TrainingSample>>();
+
+        let batcher = ChessBatcher {};
+
+        let batch = batcher.batch(samples, &device);
+
+        let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
+
+        let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
+        let p_loss_val = output.policy_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
+        let v_loss_val = output.value_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
+
+        println!(
+            "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
+            iter, loss_val, p_loss_val, v_loss_val
+        );
+
+        let grads = output.loss.backward();
+        let grads = GradientsParams::from_grads(grads, &model);
+
+        model = optim.step(training_config.lr, model, grads);
+
+        iter += 1;
+
+        if iter % 100 == 0 {
+            println!("Completed {} iterations", iter);
+        }
+
+        if training_config.max_time_s.is_some()
+            && start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
+        {
+            println!("Training stopping due to time limit...");
+            break;
+        }
+
+        if training_config.num_iters.is_some() && iter >= training_config.num_iters.unwrap() {
+            println!(
+                "Training stopping due to iteration limit ({} iters completed)...",
+                iter
+            );
+            break;
+        }
+    }
+
+    println!("Saving model...");
+
+    let metadata = ModelMetadata {
+        name: training_config.model_name,
+        iterations: iter,
+    };
+
+    std::fs::write(
+        metadata_path,
+        serde_json::to_string_pretty(&metadata)
+            .expect("Should be able to convert metadata to JSON string"),
+    )
+    .expect("Should be able to write metadata");
+
+    model
+        .into_record()
+        .save(&model_path)
+        .expect("Should be able to save the model");
+    println!("Model saved in {:?}, exiting training.", model_path);
+    return;
+}
+
+fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
+    if visits.is_empty() {
+        return Vec::new();
+    }
+
+    // Special case: deterministic selection
+    if temperature == 0.0 {
+        let (&best_idx, _) = visits
+            .iter()
+            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
+            .map(|(i, p)| (i, p))
+            .unwrap();
+
+        return vec![(best_idx, 1.0)];
+    }
+
+    let inv_temp = 1.0 / temperature;
+
+    // Step 1: apply exponent
+    let mut adjusted: Vec<(usize, f32)> =
+        visits.iter().map(|(i, v)| (*i, v.powf(inv_temp))).collect();
+
+    // Step 2: normalize
+    let sum: f32 = adjusted.iter().map(|(_, v)| *v).sum();
+
+    if sum <= 0.0 {
+        return adjusted; // fallback (shouldn't happen in normal MCTS)
+    }
+
+    for (_, v) in adjusted.iter_mut() {
+        *v /= sum;
+    }
+
+    adjusted
+}
+
+fn sample_move(
+    dist: &[(usize, f32)],
+    rng: &mut SmallRng,
+    side_to_move: Color,
+) -> Option<ChessMove> {
+    let mut r: f32 = rng.random_range(0.0..1.0);
+
+    for (idx, p) in dist {
+        r -= *p;
+        if r <= 0.0 {
+            return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
+        }
+    }
+
+    // fallback due to floating point drift
+    dist.get(0)
+        .map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
+}
diff --git a/engine/src/training.rs b/engine/src/training.rs
deleted file mode 100644
index 48b7c99..0000000
--- a/engine/src/training.rs
+++ /dev/null
@@ -1 +0,0 @@
-pub mod train;
diff --git a/engine/src/training/train.rs b/engine/src/training/train.rs
deleted file mode 100644
index faadde2..0000000
--- a/engine/src/training/train.rs
+++ /dev/null
@@ -1,292 +0,0 @@
-use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
-use crate::net::encoding::decode_move;
-use crate::net::model::{
-    ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
-};
-use burn::data::dataloader::batcher::Batcher;
-use burn::module::{AutodiffModule, Module};
-use burn::optim::{AdamConfig, GradientsParams, Optimizer};
-use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
-use burn::tensor::backend::AutodiffBackend;
-use chess::{ChessMove, Color};
-use rand::rngs::SmallRng;
-use rand::seq::SliceRandom;
-use rand::{RngExt, SeedableRng};
-use std::collections::VecDeque;
-use std::fs::File;
-use std::io::BufReader;
-use std::marker::PhantomData;
-use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
-use std::time::Instant;
-use std::time::{SystemTime, UNIX_EPOCH};
-
-pub struct TrainingConfig {
-    pub max_time_s: Option<u64>,
-    pub num_iters: Option<usize>,
-    pub max_depth: u16, // unused
-    pub model_name: String,
-    pub load_model: bool,
-    pub hidden_channels: usize,
-    pub num_blocks: usize,
-    pub batch_size: usize,      // num positions sampled to a batch (2048)
-    pub num_episodes: usize,    // num games generated per iteration (5000)
-    pub buffer_max_size: usize, // max number of samples in the buffer
-    pub mcts_config: MctsConfig,
-    pub optimizer: AdamConfig,
-    pub lr: f64,
-    pub seed: Option<u64>,
-}
-
-pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Device) {
-    let model_path = format!("artifacts/{}", training_config.model_name.as_str());
-    let metadata_path = format!("{}.json", model_path);
-    println!("Creating model...");
-    let mut model: ChessModel<B> = ChessModelConfig::init(
-        training_config.num_blocks,
-        training_config.hidden_channels,
-        &device,
-    );
-    if training_config.load_model {
-        println!("Loading model {}...", model_path);
-        // Load model in full precision from MessagePack file
-        let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
-        model = model
-            .load_file(&model_path, &recorder, &device)
-            .expect("Should be able to load the model weights from the provided file");
-    }
-
-    let train = Arc::new(AtomicBool::new(true));
-    let train_signal = Arc::clone(&train);
-
-    let mut iter: usize = 0;
-
-    if training_config.load_model {
-        let file = File::open(&metadata_path).unwrap();
-
-        // 2. Use a buffered reader for efficiency
-        let reader = BufReader::new(file);
-
-        // 3. Deserialize JSON directly from the file reader
-        let metadata: ModelMetadata = serde_json::from_reader(reader).unwrap();
-        iter = metadata.iterations;
-        println!("Loaded model had {} iters", iter);
-    }
-
-    let start_time = Instant::now();
-
-    ctrlc::set_handler(move || {
-        train_signal.store(false, Ordering::Relaxed);
-        println!("Finishing batch before exiting...");
-    })
-    .expect("Error setting Ctrl-C handler");
-
-    let mut replay_buffer: VecDeque<TrainingSample> = VecDeque::new();
-
-    let mut mcts: Mcts<B::InnerBackend> = Mcts {
-        config: training_config.mcts_config,
-        _marker: PhantomData,
-    };
-
-    // Create RNG once and reuse it for sampling and shuffling
-    // Seed from system time (platform default entropy may be unavailable in some contexts)
-    let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
-    let seed = training_config.seed.unwrap_or(now.as_nanos() as u64);
-    println!("using seed: {}", seed);
-    let mut rng = SmallRng::seed_from_u64(seed);
-
-    // Initialize optimizer once so state (moments) persist across steps
-    let mut optim = training_config.optimizer.init();
-
-    println!("Starting training...");
-    while train.load(Ordering::Relaxed) {
-        println!("Iteration: {}", iter);
-        let infer_model = model.valid();
-        // Gen samples
-        println!("Generating {} games...", training_config.num_episodes);
-        for episode in 0..training_config.num_episodes {
-            println!("Episode: {}", episode);
-            let mut board_state = BoardState::default();
-            let mut episode_buffer: Vec<MctsResults> = vec![];
-
-            let mut game_hist = String::new();
-
-            while board_state.status == BoardStateStatus::Ongoing {
-                // let before = Instant::now();
-                let results = mcts.search(&board_state, &infer_model, &device);
-                // println!("{}", before.elapsed().as_millis());
-                episode_buffer.push(results);
-
-                let temp = if board_state.halfmove_clock < 30 {
-                    1.0
-                } else {
-                    0.0
-                };
-
-                let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
-
-                let stm = board_state.board.side_to_move();
-                let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
-                // println!("playing move: {}", mv);
-                if episode == 0 {
-                    game_hist.push_str(mv.to_string().as_str());
-                    game_hist.push(' ');
-                }
-
-                board_state.apply_move(mv)
-            }
-            if episode == 0 {
-                println!("Game history of first game of iteration: {}", game_hist);
-            }
-
-            for result in episode_buffer.iter().enumerate() {
-                if board_state.status == BoardStateStatus::Stalemate
-                    || board_state.status == BoardStateStatus::Threefold
-                    || board_state.status == BoardStateStatus::FiftyMove
-                {
-                    replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
-                        result.1.clone(),
-                        0.0,
-                    ));
-                } else if board_state.status == BoardStateStatus::WhiteWinner {
-                    replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
-                        result.1.clone(),
-                        ((result.0 % 2) as f32 * -2.0) + 1.0,
-                    ));
-                } else if board_state.status == BoardStateStatus::BlackWinner {
-                    replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
-                        result.1.clone(),
-                        ((result.0 % 2) as f32 * 2.0) - 1.0,
-                    ));
-                }
-                if replay_buffer.len() > training_config.buffer_max_size {
-                    replay_buffer.pop_front();
-                }
-            }
-        }
-        // train
-
-        println!(
-            "Finished! Training on {} samples...",
-            training_config.batch_size
-        );
-        let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
-        indices.shuffle(&mut rng);
-
-        let samples = indices
-            .into_iter()
-            .take(training_config.batch_size)
-            .map(|i| replay_buffer[i].clone())
-            .collect::<Vec<TrainingSample>>();
-
-        let batcher = ChessBatcher {};
-
-        let batch = batcher.batch(samples, &device);
-
-        let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
-
-        let grads = output.loss.backward();
-        let grads = GradientsParams::from_grads(grads, &model);
-
-        model = optim.step(training_config.lr, model, grads);
-
-        iter += 1;
-
-        if iter % 100 == 0 {
-            println!("Completed {} iterations", iter);
-        }
-
-        if training_config.max_time_s.is_some()
-            && start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
-        {
-            println!("Training stopping due to time limit...");
-            break;
-        }
-
-        if training_config.num_iters.is_some() && iter >= training_config.num_iters.unwrap() {
-            println!(
-                "Training stopping due to iteration limit ({} iters completed)...",
-                iter
-            );
-            break;
-        }
-    }
-
-    println!("Saving model...");
-
-    let metadata = ModelMetadata {
-        name: training_config.model_name,
-        iterations: iter,
-    };
-
-    std::fs::write(
-        metadata_path,
-        serde_json::to_string_pretty(&metadata)
-            .expect("Should be able to convert metadata to JSON string"),
-    )
-    .expect("Should be able to write metadata");
-
-    // Save model in MessagePack format with full precision
-    let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
-    model
-        .save_file(&model_path, &recorder)
-        .expect("Should be able to save the model");
-    println!("Model saved in {:?}, exiting training.", model_path);
-    return;
-}
-
-fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
-    if visits.is_empty() {
-        return Vec::new();
-    }
-
-    // Special case: deterministic selection
-    if temperature == 0.0 {
-        let (&best_idx, _) = visits
-            .iter()
-            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
-            .map(|(i, p)| (i, p))
-            .unwrap();
-
-        return vec![(best_idx, 1.0)];
-    }
-
-    let inv_temp = 1.0 / temperature;
-
-    // Step 1: apply exponent
-    let mut adjusted: Vec<(usize, f32)> =
-        visits.iter().map(|(i, v)| (*i, v.powf(inv_temp))).collect();
-
-    // Step 2: normalize
-    let sum: f32 = adjusted.iter().map(|(_, v)| *v).sum();
-
-    if sum <= 0.0 {
-        return adjusted; // fallback (shouldn't happen in normal MCTS)
-    }
-
-    for (_, v) in adjusted.iter_mut() {
-        *v /= sum;
-    }
-
-    adjusted
-}
-
-fn sample_move(
-    dist: &[(usize, f32)],
-    rng: &mut SmallRng,
-    side_to_move: Color,
-) -> Option<ChessMove> {
-    let mut r: f32 = rng.random_range(0.0..1.0);
-
-    for (idx, p) in dist {
-        r -= *p;
-        if r <= 0.0 {
-            return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
-        }
-    }
-
-    // fallback due to floating point drift
-    dist.get(0)
-        .map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
-}
-