From 1b978b9bd85c2c50c0f4d3b44d2dd68fed42d952 Mon Sep 17 00:00:00 2001
From: Arthur Brussee <arthur.brussee@gmail.com>
Date: Thu, 5 Dec 2024 15:26:58 +0000
Subject: [PATCH] Training improvements (#56)

* Implement adam surgery

* More tweaks to match gsplat

* Fix names of settings

* cargo update

* Typo

* More tweaks to resemble gsplat

* Add radii screen splitting, fix quaternion * vec multiply

* Version bump
---
 Cargo.lock                                    | 1245 ++++++-----------
 Cargo.toml                                    |    9 +-
 crates/brush-dataset/src/lib.rs               |    4 +-
 crates/brush-dataset/src/scene_loader.rs      |   11 +-
 crates/brush-render/src/burn_glue.rs          |   21 +-
 crates/brush-render/src/gaussian_splats.rs    |   11 +-
 crates/brush-render/src/lib.rs                |    6 +-
 crates/brush-render/src/render.rs             |   26 +-
 .../src/shaders/gather_grads.wgsl             |   10 +-
 .../src/shaders/project_forward.wgsl          |    8 +-
 crates/brush-train/Cargo.toml                 |    1 +
 crates/brush-train/src/scene.rs               |    8 +-
 crates/brush-train/src/train.rs               |  557 ++++----
 crates/brush-viewer/src/panels/load_data.rs   |   25 +-
 crates/brush-viewer/src/train_loop.rs         |    8 +-
 crates/brush-viewer/src/viewer.rs             |    5 +-
 crates/train-2d/src/main.rs                   |    4 +-
 deny.toml                                     |   27 +-
 18 files changed, 851 insertions(+), 1135 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 565fd449..70821085 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "ab_glyph"
@@ -61,6 +61,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011"
 dependencies = [
  "cfg-if",
+ "const-random",
  "getrandom",
  "once_cell",
  "serde",
@@ -116,6 +117,12 @@ version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fc7eb209b1518d6bb87b283c20095f5228ecda460da70b44f0802523dea6da04"
 
+[[package]]
+name = "android-tzdata"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e999941b234f3131b00bc13c22d06e8c5ff726d1b6318ac7eb276997bbb4fef0"
+
 [[package]]
 name = "android_log-sys"
 version = "0.3.1"
@@ -150,9 +157,9 @@ checksum = "55cc3b69f167a1ef2e161439aa98aed94e6028e5f9a59be9a6ffb47aef1651f9"
 
 [[package]]
 name = "anyhow"
-version = "1.0.93"
+version = "1.0.94"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4c95c10ba0b00a02636238b814946408b1322d5ac4760326e6fb8ec956d85775"
+checksum = "c1fd03a028ef38ba2276dce7e33fcd6369c158a1bca17946c4b1b701891c1ff7"
 
 [[package]]
 name = "arbitrary"
@@ -214,254 +221,225 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
 
 [[package]]
-name = "arrow-format"
-version = "0.8.1"
+name = "arrow"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "07884ea216994cdc32a2d5f8274a8bee979cfe90274b83f86f440866ee3132c7"
+checksum = "c91839b07e474b3995035fd8ac33ee54f9c9ccbbb1ea33d9909c71bffdf1259d"
 dependencies = [
- "planus",
- "serde",
+ "arrow-arith",
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-cast",
+ "arrow-data",
+ "arrow-ord",
+ "arrow-row",
+ "arrow-schema",
+ "arrow-select",
+ "arrow-string",
 ]
 
 [[package]]
-name = "as-raw-xcb-connection"
-version = "1.0.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "175571dd1d178ced59193a6fc02dde1b972eb0bc56c892cde9beeceac5bf0f6b"
-
-[[package]]
-name = "ash"
-version = "0.38.0+1.3.281"
+name = "arrow-arith"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0bb44936d800fea8f016d7f2311c6a4f97aebd5dc86f09906139ec848cf3a46f"
+checksum = "855c57c4efd26722b044dcd3e348252560e3e0333087fb9f6479dc0bf744054f"
 dependencies = [
- "libloading",
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-data",
+ "arrow-schema",
+ "chrono",
+ "half",
+ "num",
 ]
 
 [[package]]
-name = "ashpd"
-version = "0.6.8"
+name = "arrow-array"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ac22eda5891cc086690cb6fa10121c0390de0e3b04eb269f2d766b00d3f2d81"
+checksum = "bd03279cea46569acf9295f6224fbc370c5df184b4d2ecfe97ccb131d5615a7f"
 dependencies = [
- "async-fs 2.1.2",
- "async-net",
- "enumflags2",
- "futures-channel",
- "futures-util",
- "once_cell",
- "rand",
- "serde",
- "serde_repr",
- "url",
- "zbus",
+ "ahash",
+ "arrow-buffer",
+ "arrow-data",
+ "arrow-schema",
+ "chrono",
+ "half",
+ "hashbrown 0.15.2",
+ "num",
 ]
 
 [[package]]
-name = "assert_approx_eq"
-version = "1.1.0"
+name = "arrow-buffer"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3c07dab4369547dbe5114677b33fbbf724971019f3818172d59a97a61c774ffd"
-
-[[package]]
-name = "async-broadcast"
-version = "0.5.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7c48ccdbf6ca6b121e0f586cbc0e73ae440e56c67c30fa0873b4e110d9c26d2b"
+checksum = "9e4a9b9b1d6d7117f6138e13bc4dd5daa7f94e671b70e8c9c4dc37b4f5ecfc16"
 dependencies = [
- "event-listener 2.5.3",
- "futures-core",
+ "bytes",
+ "half",
+ "num",
 ]
 
 [[package]]
-name = "async-channel"
-version = "2.3.1"
+name = "arrow-cast"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "89b47800b0be77592da0afd425cc03468052844aff33b84e33cc696f64e77b6a"
+checksum = "bc70e39916e60c5b7af7a8e2719e3ae589326039e1e863675a008bee5ffe90fd"
 dependencies = [
- "concurrent-queue",
- "event-listener-strategy",
- "futures-core",
- "pin-project-lite",
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-data",
+ "arrow-schema",
+ "arrow-select",
+ "atoi",
+ "base64",
+ "chrono",
+ "half",
+ "lexical-core",
+ "num",
+ "ryu",
 ]
 
 [[package]]
-name = "async-executor"
-version = "1.13.1"
+name = "arrow-data"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "30ca9a001c1e8ba5149f91a74362376cc6bc5b919d92d988668657bd570bdcec"
+checksum = "e4e75edf21ffd53744a9b8e3ed11101f610e7ceb1a29860432824f1834a1f623"
 dependencies = [
- "async-task",
- "concurrent-queue",
- "fastrand 2.2.0",
- "futures-lite 2.5.0",
- "slab",
+ "arrow-buffer",
+ "arrow-schema",
+ "half",
+ "num",
 ]
 
 [[package]]
-name = "async-fn-stream"
-version = "0.2.2"
+name = "arrow-format"
+version = "0.8.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e71711442f1016c768c259bec59300a10efe753bc3e686ec19e2c6a54a97c29b"
+checksum = "07884ea216994cdc32a2d5f8274a8bee979cfe90274b83f86f440866ee3132c7"
 dependencies = [
- "futures-util",
- "pin-project-lite",
+ "planus",
+ "serde",
 ]
 
 [[package]]
-name = "async-fs"
-version = "1.6.0"
+name = "arrow-ord"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "279cf904654eeebfa37ac9bb1598880884924aab82e290aa65c9e77a0e142e06"
+checksum = "ece7b5bc1180e6d82d1a60e1688c199829e8842e38497563c3ab6ea813e527fd"
 dependencies = [
- "async-lock 2.8.0",
- "autocfg",
- "blocking",
- "futures-lite 1.13.0",
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-data",
+ "arrow-schema",
+ "arrow-select",
+ "half",
+ "num",
 ]
 
 [[package]]
-name = "async-fs"
-version = "2.1.2"
+name = "arrow-row"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ebcd09b382f40fcd159c2d695175b2ae620ffa5f3bd6f664131efff4e8b9e04a"
+checksum = "745c114c8f0e8ce211c83389270de6fbe96a9088a7b32c2a041258a443fe83ff"
 dependencies = [
- "async-lock 3.4.0",
- "blocking",
- "futures-lite 2.5.0",
+ "ahash",
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-data",
+ "arrow-schema",
+ "half",
 ]
 
 [[package]]
-name = "async-io"
-version = "1.13.0"
+name = "arrow-schema"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0fc5b45d93ef0529756f812ca52e44c221b35341892d3dcc34132ac02f3dd2af"
-dependencies = [
- "async-lock 2.8.0",
- "autocfg",
- "cfg-if",
- "concurrent-queue",
- "futures-lite 1.13.0",
- "log",
- "parking",
- "polling 2.8.0",
- "rustix 0.37.27",
- "slab",
- "socket2 0.4.10",
- "waker-fn",
-]
+checksum = "b95513080e728e4cec37f1ff5af4f12c9688d47795d17cda80b6ec2cf74d4678"
 
 [[package]]
-name = "async-io"
-version = "2.4.0"
+name = "arrow-select"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "43a2b323ccce0a1d90b449fd71f2a06ca7faa7c54c2751f06c9bd851fc061059"
+checksum = "8e415279094ea70323c032c6e739c48ad8d80e78a09bef7117b8718ad5bf3722"
 dependencies = [
- "async-lock 3.4.0",
- "cfg-if",
- "concurrent-queue",
- "futures-io",
- "futures-lite 2.5.0",
- "parking",
- "polling 3.7.4",
- "rustix 0.38.41",
- "slab",
- "tracing",
- "windows-sys 0.59.0",
+ "ahash",
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-data",
+ "arrow-schema",
+ "num",
 ]
 
 [[package]]
-name = "async-lock"
-version = "2.8.0"
+name = "arrow-string"
+version = "53.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "287272293e9d8c41773cec55e365490fe034813a2f172f502d6ddcf75b2f582b"
+checksum = "11d956cae7002eb8d83a27dbd34daaea1cf5b75852f0b84deb4d93a276e92bbf"
 dependencies = [
- "event-listener 2.5.3",
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-data",
+ "arrow-schema",
+ "arrow-select",
+ "memchr",
+ "num",
+ "regex",
+ "regex-syntax 0.8.5",
 ]
 
 [[package]]
-name = "async-lock"
-version = "3.4.0"
+name = "as-raw-xcb-connection"
+version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ff6e472cdea888a4bd64f342f09b3f50e1886d32afe8df3d663c01140b811b18"
-dependencies = [
- "event-listener 5.3.1",
- "event-listener-strategy",
- "pin-project-lite",
-]
+checksum = "175571dd1d178ced59193a6fc02dde1b972eb0bc56c892cde9beeceac5bf0f6b"
 
 [[package]]
-name = "async-net"
-version = "2.0.0"
+name = "ash"
+version = "0.38.0+1.3.281"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b948000fad4873c1c9339d60f2623323a0cfd3816e5181033c6a5cb68b2accf7"
+checksum = "0bb44936d800fea8f016d7f2311c6a4f97aebd5dc86f09906139ec848cf3a46f"
 dependencies = [
- "async-io 2.4.0",
- "blocking",
- "futures-lite 2.5.0",
+ "libloading",
 ]
 
 [[package]]
-name = "async-process"
-version = "1.8.1"
+name = "assert_approx_eq"
+version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ea6438ba0a08d81529c69b36700fa2f95837bfe3e776ab39cde9c14d9149da88"
-dependencies = [
- "async-io 1.13.0",
- "async-lock 2.8.0",
- "async-signal",
- "blocking",
- "cfg-if",
- "event-listener 3.1.0",
- "futures-lite 1.13.0",
- "rustix 0.38.41",
- "windows-sys 0.48.0",
-]
+checksum = "3c07dab4369547dbe5114677b33fbbf724971019f3818172d59a97a61c774ffd"
 
 [[package]]
-name = "async-recursion"
-version = "1.1.1"
+name = "async-channel"
+version = "2.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3b43422f69d8ff38f95f1b2bb76517c91589a924d1559a0e935d7c8ce0274c11"
+checksum = "89b47800b0be77592da0afd425cc03468052844aff33b84e33cc696f64e77b6a"
 dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.90",
+ "concurrent-queue",
+ "event-listener-strategy",
+ "futures-core",
+ "pin-project-lite",
 ]
 
 [[package]]
-name = "async-signal"
-version = "0.2.10"
+name = "async-fn-stream"
+version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "637e00349800c0bdf8bfc21ebbc0b6524abea702b0da4168ac00d070d0c0b9f3"
+checksum = "e71711442f1016c768c259bec59300a10efe753bc3e686ec19e2c6a54a97c29b"
 dependencies = [
- "async-io 2.4.0",
- "async-lock 3.4.0",
- "atomic-waker",
- "cfg-if",
- "futures-core",
- "futures-io",
- "rustix 0.38.41",
- "signal-hook-registry",
- "slab",
- "windows-sys 0.59.0",
+ "futures-util",
+ "pin-project-lite",
 ]
 
 [[package]]
-name = "async-task"
-version = "4.7.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de"
-
-[[package]]
-name = "async-trait"
-version = "0.1.83"
+name = "async-lock"
+version = "3.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "721cae7de5c34fbb2acd27e21e6d2cf7b886dce0c27388d46c4e6c47ea4318dd"
+checksum = "ff6e472cdea888a4bd64f342f09b3f50e1886d32afe8df3d663c01140b811b18"
 dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.90",
+ "event-listener",
+ "event-listener-strategy",
+ "pin-project-lite",
 ]
 
 [[package]]
@@ -476,6 +454,15 @@ dependencies = [
  "system-deps",
 ]
 
+[[package]]
+name = "atoi"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528"
+dependencies = [
+ "num-traits",
+]
+
 [[package]]
 name = "atomic-waker"
 version = "1.1.2"
@@ -559,15 +546,6 @@ version = "1.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b"
 
-[[package]]
-name = "bincode"
-version = "1.3.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad"
-dependencies = [
- "serde",
-]
-
 [[package]]
 name = "bincode"
 version = "2.0.0-rc.3"
@@ -655,19 +633,6 @@ dependencies = [
  "objc2",
 ]
 
-[[package]]
-name = "blocking"
-version = "1.6.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "703f41c54fc768e63e091340b424302bb1c29ef4aa0c7f10fe849dfb114d29ea"
-dependencies = [
- "async-channel",
- "async-task",
- "futures-io",
- "futures-lite 2.5.0",
- "piper",
-]
-
 [[package]]
 name = "brush-android"
 version = "0.1.0"
@@ -825,6 +790,7 @@ dependencies = [
  "brush-render",
  "burn",
  "glam",
+ "hashbrown 0.15.2",
  "image",
  "log",
  "rand",
@@ -892,7 +858,7 @@ dependencies = [
  "naga_oil",
  "regex",
  "regex-syntax 0.8.5",
- "thiserror 2.0.3",
+ "thiserror 2.0.4",
 ]
 
 [[package]]
@@ -921,7 +887,7 @@ checksum = "79296716171880943b8470b5f8d03aa55eb2e645a4874bdbb28adb49162e012c"
 [[package]]
 name = "burn"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-core",
  "burn-train",
@@ -930,7 +896,7 @@ dependencies = [
 [[package]]
 name = "burn-autodiff"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-common",
  "burn-tensor",
@@ -942,7 +908,7 @@ dependencies = [
 [[package]]
 name = "burn-candle"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-tensor",
  "candle-core",
@@ -953,7 +919,7 @@ dependencies = [
 [[package]]
 name = "burn-common"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "cubecl-common",
  "getrandom",
@@ -965,10 +931,10 @@ dependencies = [
 [[package]]
 name = "burn-core"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "ahash",
- "bincode 2.0.0-rc.3",
+ "bincode",
  "burn-autodiff",
  "burn-candle",
  "burn-common",
@@ -1000,7 +966,7 @@ dependencies = [
 [[package]]
 name = "burn-cuda"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-fusion",
  "burn-jit",
@@ -1015,7 +981,7 @@ dependencies = [
 [[package]]
 name = "burn-dataset"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "csv",
  "derive-new 0.7.0",
@@ -1034,13 +1000,13 @@ dependencies = [
  "strum",
  "strum_macros",
  "tempfile",
- "thiserror 2.0.3",
+ "thiserror 2.0.4",
 ]
 
 [[package]]
 name = "burn-derive"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "derive-new 0.7.0",
  "proc-macro2",
@@ -1051,7 +1017,7 @@ dependencies = [
 [[package]]
 name = "burn-fusion"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-common",
  "burn-tensor",
@@ -1066,7 +1032,7 @@ dependencies = [
 [[package]]
 name = "burn-hip"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-fusion",
  "burn-jit",
@@ -1081,7 +1047,7 @@ dependencies = [
 [[package]]
 name = "burn-jit"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-common",
  "burn-fusion",
@@ -1089,7 +1055,7 @@ dependencies = [
  "bytemuck",
  "cubecl",
  "derive-new 0.7.0",
- "futures-lite 2.5.0",
+ "futures-lite",
  "half",
  "hashbrown 0.15.2",
  "log",
@@ -1103,7 +1069,7 @@ dependencies = [
 [[package]]
 name = "burn-ndarray"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "atomic_float",
  "burn-autodiff",
@@ -1122,7 +1088,7 @@ dependencies = [
 [[package]]
 name = "burn-router"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-common",
  "burn-tensor",
@@ -1134,7 +1100,7 @@ dependencies = [
 [[package]]
 name = "burn-tch"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-tensor",
  "half",
@@ -1147,7 +1113,7 @@ dependencies = [
 [[package]]
 name = "burn-tensor"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-common",
  "bytemuck",
@@ -1167,7 +1133,7 @@ dependencies = [
 [[package]]
 name = "burn-train"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "async-channel",
  "burn-core",
@@ -1188,7 +1154,7 @@ dependencies = [
 [[package]]
 name = "burn-wgpu"
 version = "0.16.0"
-source = "git+https://github.com/tracel-ai/burn#3dc4b43e92ec5683ca39c194acf274b131d44dfa"
+source = "git+https://github.com/tracel-ai/burn#6dbc3fd295c8c84009f32c012f4161bd81106b73"
 dependencies = [
  "burn-fusion",
  "burn-jit",
@@ -1285,8 +1251,8 @@ checksum = "b99da2f8558ca23c71f4fd15dc57c906239752dd27ff3c00a1d56b685b7cbfec"
 dependencies = [
  "bitflags 2.6.0",
  "log",
- "polling 3.7.4",
- "rustix 0.38.41",
+ "polling",
+ "rustix",
  "slab",
  "thiserror 1.0.69",
 ]
@@ -1298,7 +1264,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "95a66a987056935f7efce4ab5668920b5d0dac4a7c99991a67395f13702ddd20"
 dependencies = [
  "calloop",
- "rustix 0.38.41",
+ "rustix",
  "wayland-backend",
  "wayland-client",
 ]
@@ -1447,7 +1413,10 @@ version = "0.4.38"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a21f936df1771bf62b77f047b726c4625ff2e8aa607c01ec06e5a05bd8463401"
 dependencies = [
+ "android-tzdata",
+ "iana-time-zone",
  "num-traits",
+ "windows-targets 0.52.6",
 ]
 
 [[package]]
@@ -1471,18 +1440,18 @@ dependencies = [
 
 [[package]]
 name = "clap"
-version = "4.5.21"
+version = "4.5.22"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fb3b4b9e5a7c7514dfa52869339ee98b3156b0bfb4e8a77c4ff4babb64b1604f"
+checksum = "69371e34337c4c984bbe322360c2547210bf632eb2814bbe78a6e87a2935bd2b"
 dependencies = [
  "clap_builder",
 ]
 
 [[package]]
 name = "clap_builder"
-version = "4.5.21"
+version = "4.5.22"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b17a95aa67cc7b5ebd32aa5370189aa0d79069ef1c64ce893bd30fb24bff20ec"
+checksum = "6e24c1b4099818523236a8ca881d2b45db98dadfb4625cf6608c12069fcbbde1"
 dependencies = [
  "anstyle",
  "clap_lex",
@@ -1648,6 +1617,26 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "const-random"
+version = "0.1.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359"
+dependencies = [
+ "const-random-macro",
+]
+
+[[package]]
+name = "const-random-macro"
+version = "0.1.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e"
+dependencies = [
+ "getrandom",
+ "once_cell",
+ "tiny-keccak",
+]
+
 [[package]]
 name = "constant_time_eq"
 version = "0.1.5"
@@ -1797,7 +1786,7 @@ dependencies = [
  "crossterm_winapi",
  "mio",
  "parking_lot",
- "rustix 0.38.41",
+ "rustix",
  "signal-hook",
  "signal-hook-mio",
  "winapi",
@@ -1852,7 +1841,7 @@ dependencies = [
 [[package]]
 name = "cubecl"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "cubecl-core",
  "cubecl-cuda",
@@ -1866,11 +1855,11 @@ dependencies = [
 [[package]]
 name = "cubecl-common"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "derive-new 0.6.0",
  "embassy-futures",
- "futures-lite 2.5.0",
+ "futures-lite",
  "getrandom",
  "log",
  "portable-atomic",
@@ -1883,7 +1872,7 @@ dependencies = [
 [[package]]
 name = "cubecl-core"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1901,7 +1890,7 @@ dependencies = [
 [[package]]
 name = "cubecl-cpp"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1915,7 +1904,7 @@ dependencies = [
 [[package]]
 name = "cubecl-cuda"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1931,7 +1920,7 @@ dependencies = [
 [[package]]
 name = "cubecl-hip"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "bytemuck",
  "cubecl-common",
@@ -1957,7 +1946,7 @@ dependencies = [
 [[package]]
 name = "cubecl-linalg"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "bytemuck",
  "cubecl-core",
@@ -1968,7 +1957,7 @@ dependencies = [
 [[package]]
 name = "cubecl-macros"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "cubecl-common",
  "darling",
@@ -1983,10 +1972,10 @@ dependencies = [
 [[package]]
 name = "cubecl-runtime"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "async-channel",
- "async-lock 3.4.0",
+ "async-lock",
  "cfg_aliases 0.2.1",
  "cubecl-common",
  "derive-new 0.6.0",
@@ -2004,7 +1993,7 @@ dependencies = [
 [[package]]
 name = "cubecl-wgpu"
 version = "0.4.0"
-source = "git+https://github.com/tracel-ai/cubecl?rev=a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6#a4e2b77dcc1c91e31ca95a8d55454f8c49e1f4f6"
+source = "git+https://github.com/tracel-ai/cubecl?rev=1c4e0036c043422773fd6824c2a888160fca8e5e#1c4e0036c043422773fd6824c2a888160fca8e5e"
 dependencies = [
  "async-channel",
  "bytemuck",
@@ -2109,17 +2098,6 @@ dependencies = [
  "powerfmt",
 ]
 
-[[package]]
-name = "derivative"
-version = "2.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fcc3dd5e9e9c0b295d6e1e4d811fb6f157d5ffd784b8d202fc62eac8035a770b"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 1.0.109",
-]
-
 [[package]]
 name = "derive-new"
 version = "0.6.0"
@@ -2231,9 +2209,9 @@ dependencies = [
 
 [[package]]
 name = "divan"
-version = "0.1.16"
+version = "0.1.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ccc40f214f0d9e897cfc72e2edfa5c225d3252f758c537f11ac0a80371c073a6"
+checksum = "e0583193020b29b03682d8d33bb53a5b0f50df6daacece12ca99b904cfdcb8c4"
 dependencies = [
  "cfg-if",
  "clap",
@@ -2245,9 +2223,9 @@ dependencies = [
 
 [[package]]
 name = "divan-macros"
-version = "0.1.16"
+version = "0.1.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7bdb5411188f7f878a17964798c1264b6b0a9f915bd39b20bf99193c923e1b4e"
+checksum = "8dc51d98e636f5e3b0759a39257458b22619cac7e96d932da6eeb052891bb67c"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -2312,22 +2290,13 @@ dependencies = [
  "reborrow",
 ]
 
-[[package]]
-name = "ecolor"
-version = "0.29.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "775cfde491852059e386c4e1deb4aef381c617dc364184c6f6afee99b87c402b"
-dependencies = [
- "emath 0.29.1 (registry+https://github.com/rust-lang/crates.io-index)",
-]
-
 [[package]]
 name = "ecolor"
 version = "0.29.1"
 source = "git+https://github.com/emilk/egui/?rev=5bfff316c9818b3c140d02bb6cdc488556d46ab7#5bfff316c9818b3c140d02bb6cdc488556d46ab7"
 dependencies = [
  "bytemuck",
- "emath 0.29.1 (git+https://github.com/emilk/egui/?rev=5bfff316c9818b3c140d02bb6cdc488556d46ab7)",
+ "emath",
  "serde",
 ]
 
@@ -2377,7 +2346,7 @@ source = "git+https://github.com/emilk/egui/?rev=5bfff316c9818b3c140d02bb6cdc488
 dependencies = [
  "accesskit",
  "ahash",
- "emath 0.29.1 (git+https://github.com/emilk/egui/?rev=5bfff316c9818b3c140d02bb6cdc488556d46ab7)",
+ "emath",
  "epaint",
  "log",
  "nohash-hasher",
@@ -2428,7 +2397,7 @@ dependencies = [
  "egui",
  "glow 0.14.2",
  "log",
- "memoffset 0.9.1",
+ "memoffset",
  "wasm-bindgen",
  "web-sys",
  "winit",
@@ -2458,12 +2427,6 @@ version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6f4e5af126dafd0741c2ad62d47f68b28602550102e5f0dd45c8a97fc8b49c29"
 
-[[package]]
-name = "emath"
-version = "0.29.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b1fe0049ce51d0fb414d029e668dd72eb30bc2b739bf34296ed97bd33df544f3"
-
 [[package]]
 name = "emath"
 version = "0.29.1"
@@ -2497,27 +2460,6 @@ dependencies = [
  "syn 2.0.90",
 ]
 
-[[package]]
-name = "enumflags2"
-version = "0.7.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d232db7f5956f3f14313dc2f87985c58bd2c695ce124c8cdd984e08e15ac133d"
-dependencies = [
- "enumflags2_derive",
- "serde",
-]
-
-[[package]]
-name = "enumflags2_derive"
-version = "0.7.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "de0d48a183585823424a4ce1aa132d174a6a81bd540895822eb4c8373a8e49e8"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.90",
-]
-
 [[package]]
 name = "enumn"
 version = "0.1.14"
@@ -2560,8 +2502,8 @@ dependencies = [
  "ab_glyph",
  "ahash",
  "bytemuck",
- "ecolor 0.29.1 (git+https://github.com/emilk/egui/?rev=5bfff316c9818b3c140d02bb6cdc488556d46ab7)",
- "emath 0.29.1 (git+https://github.com/emilk/egui/?rev=5bfff316c9818b3c140d02bb6cdc488556d46ab7)",
+ "ecolor",
+ "emath",
  "epaint_default_fonts",
  "log",
  "nohash-hasher",
@@ -2611,23 +2553,6 @@ version = "1.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b90ca2580b73ab6a1f724b76ca11ab632df820fd6040c336200d2c1df7b3c82c"
 
-[[package]]
-name = "event-listener"
-version = "2.5.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0"
-
-[[package]]
-name = "event-listener"
-version = "3.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d93877bcde0eb80ca09131a08d23f0a5c18a620b01db137dba666d18cd9b30c2"
-dependencies = [
- "concurrent-queue",
- "parking",
- "pin-project-lite",
-]
-
 [[package]]
 name = "event-listener"
 version = "5.3.1"
@@ -2645,7 +2570,7 @@ version = "0.5.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3c3e4e0dd3673c1139bf041f3008816d9cf2946bbfac2945c09e523b8d7b05b2"
 dependencies = [
- "event-listener 5.3.1",
+ "event-listener",
  "pin-project-lite",
 ]
 
@@ -2685,15 +2610,6 @@ dependencies = [
  "serde",
 ]
 
-[[package]]
-name = "fastrand"
-version = "1.9.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e51093e27b0797c359783294ca4f0a911c270184cb10f85783b118614a1501be"
-dependencies = [
- "instant",
-]
-
 [[package]]
 name = "fastrand"
 version = "2.2.0"
@@ -2709,16 +2625,6 @@ dependencies = [
  "simd-adler32",
 ]
 
-[[package]]
-name = "fixed"
-version = "0.5.7"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f32ca1abdbb21d63a3e02a658a9e3001b172f13c8b46724299e21190c5ee5041"
-dependencies = [
- "serde",
- "typenum",
-]
-
 [[package]]
 name = "fixed"
 version = "1.28.0"
@@ -2729,6 +2635,7 @@ dependencies = [
  "bytemuck",
  "half",
  "num-traits",
+ "serde",
  "typenum",
 ]
 
@@ -2854,28 +2761,13 @@ version = "0.3.31"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6"
 
-[[package]]
-name = "futures-lite"
-version = "1.13.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "49a9d51ce47660b1e808d3c990b4709f2f415d928835a17dfd16991515c46bce"
-dependencies = [
- "fastrand 1.9.0",
- "futures-core",
- "futures-io",
- "memchr",
- "parking",
- "pin-project-lite",
- "waker-fn",
-]
-
 [[package]]
 name = "futures-lite"
 version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "cef40d21ae2c515b51041df9ed313ed21e572df340ea58a922a0aefe7e8891a1"
 dependencies = [
- "fastrand 2.2.0",
+ "fastrand",
  "futures-core",
  "futures-io",
  "parking",
@@ -3079,9 +2971,9 @@ dependencies = [
 
 [[package]]
 name = "generator"
-version = "0.8.3"
+version = "0.8.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dbb949699c3e4df3a183b1d2142cb24277057055ed23c68ed58894f76c517223"
+checksum = "cc6bd114ceda131d3b1d665eba35788690ad37f5916457286b32ab6fd3c438dd"
 dependencies = [
  "cfg-if",
  "libc",
@@ -3170,7 +3062,7 @@ version = "0.12.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "34740384d8d763975858fa2c176b68652a6fcc09f616e24e3ce967b0d370e4d8"
 dependencies = [
- "fastrand 2.2.0",
+ "fastrand",
  "gix-features",
  "gix-utils",
 ]
@@ -3182,7 +3074,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0b5eccc17194ed0e67d49285e4853307e4147e95407f91c1c3e4a13ba9f4e4ce"
 dependencies = [
  "faster-hex",
- "thiserror 2.0.3",
+ "thiserror 2.0.4",
 ]
 
 [[package]]
@@ -3213,7 +3105,7 @@ version = "0.1.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ba427e3e9599508ed98a6ddf8ed05493db114564e338e41f6a996d2e4790335f"
 dependencies = [
- "fastrand 2.2.0",
+ "fastrand",
  "unicode-normalization",
 ]
 
@@ -3461,12 +3353,6 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fbf6a919d6cf397374f7dfeeea91d974c7c0a7221d0d0f4f20d859d329e53fcc"
 
-[[package]]
-name = "hex"
-version = "0.4.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
-
 [[package]]
 name = "hexf-parse"
 version = "0.2.1"
@@ -3493,9 +3379,9 @@ dependencies = [
 
 [[package]]
 name = "http"
-version = "1.1.0"
+version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "21b9ddb458710bc376481b842f5da65cdf31522de232c1ca8146abce2a358258"
+checksum = "f16ca2af56261c99fba8bac40a10251ce8188205a4c448fbb745a2e4daa76fea"
 dependencies = [
  "bytes",
  "fnv",
@@ -3587,12 +3473,35 @@ dependencies = [
  "http-body",
  "hyper",
  "pin-project-lite",
- "socket2 0.5.8",
+ "socket2",
  "tokio",
  "tower-service",
  "tracing",
 ]
 
+[[package]]
+name = "iana-time-zone"
+version = "0.1.61"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "235e081f3925a06703c2d0117ea8b91f042756fd6e7a6e5d901e8ca1a996b220"
+dependencies = [
+ "android_system_properties",
+ "core-foundation-sys",
+ "iana-time-zone-haiku",
+ "js-sys",
+ "wasm-bindgen",
+ "windows-core 0.52.0",
+]
+
+[[package]]
+name = "iana-time-zone-haiku"
+version = "0.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
+dependencies = [
+ "cc",
+]
+
 [[package]]
 name = "icu_collections"
 version = "1.5.0"
@@ -3801,9 +3710,9 @@ checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5"
 
 [[package]]
 name = "infer"
-version = "0.15.0"
+version = "0.16.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cb33622da908807a06f9513c19b3c1ad50fab3e4137d82a78107d502075aa199"
+checksum = "bc150e5ce2330295b8616ce0e3f53250e53af31759a9dbedad1621ba29151847"
 dependencies = [
  "cfb",
 ]
@@ -3837,15 +3746,6 @@ dependencies = [
  "syn 2.0.90",
 ]
 
-[[package]]
-name = "instant"
-version = "0.1.13"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222"
-dependencies = [
- "cfg-if",
-]
-
 [[package]]
 name = "interpolate_name"
 version = "0.2.4"
@@ -3857,17 +3757,6 @@ dependencies = [
  "syn 2.0.90",
 ]
 
-[[package]]
-name = "io-lifetimes"
-version = "1.0.11"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "eae7b9aee968036d54dce06cebaefd919e4472e753296daccd6d344e3e2df0c2"
-dependencies = [
- "hermit-abi 0.3.9",
- "libc",
- "windows-sys 0.48.0",
-]
-
 [[package]]
 name = "ipnet"
 version = "2.10.1"
@@ -3998,7 +3887,7 @@ dependencies = [
  "divrem",
  "doc-comment",
  "elapsed",
- "fixed 1.28.0",
+ "fixed",
  "generator",
  "init_with",
  "itertools 0.13.0",
@@ -4023,6 +3912,70 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "03087c2bad5e1034e8cace5926dec053fb3790248370865f5117a7d0213354c8"
 
+[[package]]
+name = "lexical-core"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0431c65b318a590c1de6b8fd6e72798c92291d27762d94c9e6c37ed7a73d8458"
+dependencies = [
+ "lexical-parse-float",
+ "lexical-parse-integer",
+ "lexical-util",
+ "lexical-write-float",
+ "lexical-write-integer",
+]
+
+[[package]]
+name = "lexical-parse-float"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eb17a4bdb9b418051aa59d41d65b1c9be5affab314a872e5ad7f06231fb3b4e0"
+dependencies = [
+ "lexical-parse-integer",
+ "lexical-util",
+ "static_assertions",
+]
+
+[[package]]
+name = "lexical-parse-integer"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5df98f4a4ab53bf8b175b363a34c7af608fe31f93cc1fb1bf07130622ca4ef61"
+dependencies = [
+ "lexical-util",
+ "static_assertions",
+]
+
+[[package]]
+name = "lexical-util"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85314db53332e5c192b6bca611fb10c114a80d1b831ddac0af1e9be1b9232ca0"
+dependencies = [
+ "static_assertions",
+]
+
+[[package]]
+name = "lexical-write-float"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6e7c3ad4e37db81c1cbe7cf34610340adc09c322871972f74877a712abc6c809"
+dependencies = [
+ "lexical-util",
+ "lexical-write-integer",
+ "static_assertions",
+]
+
+[[package]]
+name = "lexical-write-integer"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "eb89e9f6958b83258afa3deed90b5de9ef68eef090ad5086c791cd2345610162"
+dependencies = [
+ "lexical-util",
+ "static_assertions",
+]
+
 [[package]]
 name = "libc"
 version = "0.2.167"
@@ -4083,12 +4036,6 @@ version = "0.5.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0717cef1bc8b636c6e1c1bbdefc09e6322da8a9321966e8928ef80d20f7f770f"
 
-[[package]]
-name = "linux-raw-sys"
-version = "0.3.8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ef53942eb7bf7ff43a617b3e2c1c4a5ecf5944a7c1bc12d7ee39bbb15e5c1519"
-
 [[package]]
 name = "linux-raw-sys"
 version = "0.4.14"
@@ -4235,19 +4182,10 @@ checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
 name = "memmap2"
 version = "0.9.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f"
-dependencies = [
- "libc",
- "stable_deref_trait",
-]
-
-[[package]]
-name = "memoffset"
-version = "0.7.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5de893c32cde5f383baa4c04c5d6dbdd735cfd4a794b0debdb2bb1b421da5ff4"
+checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f"
 dependencies = [
- "autocfg",
+ "libc",
+ "stable_deref_trait",
 ]
 
 [[package]]
@@ -4480,18 +4418,6 @@ version = "1.0.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "650eef8c711430f1a879fdd01d4745a7deea475becfb90269c06775983bbf086"
 
-[[package]]
-name = "nix"
-version = "0.26.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "598beaf3cc6fdd9a5dfb1630c2800c7acd31df7aaf0f565796fba2b53ca1af1b"
-dependencies = [
- "bitflags 1.3.2",
- "cfg-if",
- "libc",
- "memoffset 0.7.1",
-]
-
 [[package]]
 name = "nohash-hasher"
 version = "0.2.0"
@@ -4651,7 +4577,7 @@ version = "0.7.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "af1844ef2428cc3e1cb900be36181049ef3d3193c63e43026cfe202983b27a56"
 dependencies = [
- "proc-macro-crate 3.2.0",
+ "proc-macro-crate",
  "proc-macro2",
  "quote",
  "syn 2.0.90",
@@ -4960,16 +4886,6 @@ dependencies = [
  "num-traits",
 ]
 
-[[package]]
-name = "ordered-stream"
-version = "0.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9aa2b01e1d916879f73a53d01d1d6cee68adbb31d6d9177a8cfce093cced1d50"
-dependencies = [
- "futures-core",
- "pin-project-lite",
-]
-
 [[package]]
 name = "overload"
 version = "0.1.1"
@@ -5153,17 +5069,6 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184"
 
-[[package]]
-name = "piper"
-version = "0.2.4"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "96c8c490f422ef9a4efd2cb5b42b76c8613d7e7dfc1caf667b8a3350a5acc066"
-dependencies = [
- "atomic-waker",
- "fastrand 2.2.0",
- "futures-io",
-]
-
 [[package]]
 name = "pkg-config"
 version = "0.3.31"
@@ -5216,22 +5121,6 @@ dependencies = [
  "miniz_oxide",
 ]
 
-[[package]]
-name = "polling"
-version = "2.8.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4b2d323e8ca7996b3e23126511a523f7e62924d93ecd5ae73b333815b0eb3dce"
-dependencies = [
- "autocfg",
- "bitflags 1.3.2",
- "cfg-if",
- "concurrent-queue",
- "libc",
- "log",
- "pin-project-lite",
- "windows-sys 0.48.0",
-]
-
 [[package]]
 name = "polling"
 version = "3.7.4"
@@ -5242,7 +5131,7 @@ dependencies = [
  "concurrent-queue",
  "hermit-abi 0.4.0",
  "pin-project-lite",
- "rustix 0.38.41",
+ "rustix",
  "tracing",
  "windows-sys 0.59.0",
 ]
@@ -5318,23 +5207,13 @@ dependencies = [
  "syn 2.0.90",
 ]
 
-[[package]]
-name = "proc-macro-crate"
-version = "1.3.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7f4c021e1093a56626774e81216a4ce732a735e5bad4868a03f3ed65ca0c3919"
-dependencies = [
- "once_cell",
- "toml_edit 0.19.15",
-]
-
 [[package]]
 name = "proc-macro-crate"
 version = "3.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8ecf48c7ca261d60b74ab1a7b20da18bede46776b2e55535cb958eb595c5fa7b"
 dependencies = [
- "toml_edit 0.22.22",
+ "toml_edit",
 ]
 
 [[package]]
@@ -5372,27 +5251,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "fa9dae7b05c02ec1a6bc9bcf20d8bc64a7dcbf57934107902a872014899b741f"
 dependencies = [
  "anyhow",
- "bincode 1.3.3",
  "byteorder",
  "cfg-if",
  "itertools 0.10.5",
- "lz4_flex",
  "once_cell",
  "parking_lot",
- "serde",
-]
-
-[[package]]
-name = "puffin_http"
-version = "0.16.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "739a3c7f56604713b553d7addd7718c226e88d598979ae3450320800bd0e9810"
-dependencies = [
- "anyhow",
- "crossbeam-channel",
- "log",
- "parking_lot",
- "puffin",
 ]
 
 [[package]]
@@ -5454,8 +5317,8 @@ dependencies = [
  "quinn-udp",
  "rustc-hash 2.1.0",
  "rustls",
- "socket2 0.5.8",
- "thiserror 2.0.3",
+ "socket2",
+ "thiserror 2.0.4",
  "tokio",
  "tracing",
 ]
@@ -5474,7 +5337,7 @@ dependencies = [
  "rustls",
  "rustls-pki-types",
  "slab",
- "thiserror 2.0.3",
+ "thiserror 2.0.4",
  "tinyvec",
  "tracing",
  "web-time",
@@ -5489,7 +5352,7 @@ dependencies = [
  "cfg_aliases 0.2.1",
  "libc",
  "once_cell",
- "socket2 0.5.8",
+ "socket2",
  "tracing",
  "windows-sys 0.59.0",
 ]
@@ -5697,7 +5560,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "787fa1df3020f018e02c1f957edfc6890a73372444de397c36011cda61c9b489"
 dependencies = [
  "ahash",
+ "arrow-array",
+ "arrow-buffer",
+ "arrow-data",
  "arrow-format",
+ "arrow-schema",
  "bytemuck",
  "chrono",
  "comfy-table",
@@ -5715,18 +5582,18 @@ dependencies = [
 
 [[package]]
 name = "re_build_info"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "dd39f5cc762cb2189de8340d998b67bd96c42b50cc13d09ffed4a6c91946d64f"
+checksum = "517bef56ed778112195a374333465ec3e4dc85fa9f490b9539be599434cbe8ba"
 dependencies = [
  "serde",
 ]
 
 [[package]]
 name = "re_build_tools"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bbaeb7bd5db56a4eb9b323eabf9acb3f0552b9f999ce8224d9b5b08547f417e3"
+checksum = "d44a6a1eb8033fd1322cef92fbc0d2f1ae70e1d72e3fd949c998fdd52167c4c9"
 dependencies = [
  "anyhow",
  "cargo_metadata 0.18.1",
@@ -5739,18 +5606,18 @@ dependencies = [
 
 [[package]]
 name = "re_case"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "25c05e42afb6ad68408ab44bb20ae9207dab8162fd6f6d0aacb46064d5ad36aa"
+checksum = "1f8d3c1ab2e48402e9db14c4ae219113d64cb6892a1da3859ca4cbaf3e0b7eae"
 dependencies = [
  "convert_case",
 ]
 
 [[package]]
 name = "re_chunk"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "835c0008466c91d037d9bed4cc22ac282217bdb2e3e8cfcee2f6eef790550304"
+checksum = "da6d64af273dde39ddb3d97804b41a0f701ff649326a3f281298fd9ec105a523"
 dependencies = [
  "ahash",
  "anyhow",
@@ -5776,9 +5643,9 @@ dependencies = [
 
 [[package]]
 name = "re_chunk_store"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9f38aad1c54fb8c93bf0dc982f891552efdbf704f29080cfcd477fe5d25ea9d0"
+checksum = "3a63ee2c9d18811c3bbbc36dfbfcf5dbacdd2fd0c28877624557fbc83afa1464"
 dependencies = [
  "ahash",
  "anyhow",
@@ -5802,9 +5669,9 @@ dependencies = [
 
 [[package]]
 name = "re_crash_handler"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "317af2231c34496ff40d9bb40875f1856381c7f4ffcdb6e60e4f9a32f6ff9298"
+checksum = "d391e5bbaabc4aaacd4c82b92df73b4f492ec37e049737e624001f9961cd4962"
 dependencies = [
  "backtrace",
  "econtext",
@@ -5814,41 +5681,21 @@ dependencies = [
  "re_build_info",
 ]
 
-[[package]]
-name = "re_dataframe"
-version = "0.19.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6bbdd427394ab73e870beb592bd24d14d43ef5fc99f1f75e282ab414d88250b6"
-dependencies = [
- "anyhow",
- "itertools 0.13.0",
- "nohash-hasher",
- "re_arrow2",
- "re_chunk",
- "re_chunk_store",
- "re_log",
- "re_log_types",
- "re_query",
- "re_tracing",
- "re_types_core",
-]
-
 [[package]]
 name = "re_entity_db"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9b34166585a6d17a61662fbc8433816f9a16632208311f9db0303f0a53133f1a"
+checksum = "89fd21cf079b3d31d728867b8ee74349bfd2332cc90cf1edf46e6b400a4bbadf"
 dependencies = [
  "ahash",
  "document-features",
- "emath 0.29.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "emath",
  "itertools 0.13.0",
  "nohash-hasher",
  "parking_lot",
  "re_build_info",
  "re_chunk",
  "re_chunk_store",
- "re_dataframe",
  "re_format",
  "re_int_histogram",
  "re_log",
@@ -5864,24 +5711,24 @@ dependencies = [
 
 [[package]]
 name = "re_error"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fb97bdbd6c379c82cb8f77f01bf4543bee5699cefe1ea78932338530f7669bd0"
+checksum = "56a532660bca60b95f2d3baeb6c0136fb61dfd2e9ce2759e4f8828b8a013107b"
 
 [[package]]
 name = "re_format"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9f0676202e29d77150df3f1e7e010b5a0d9af6f0f8d430903288832748b7df57"
+checksum = "ff9572db5ecc863e03d7d6cd5a8062edda8a4f2bfff91bc9423d132d68146b19"
 dependencies = [
  "num-traits",
 ]
 
 [[package]]
 name = "re_format_arrow"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ce2a885e055eb557923f56e52a1249c026a6d06fa4399dbc356308fed4e98058"
+checksum = "7e264d1a6593ff3fdc54ab6ed430e752c25749c4e5268b019c6b8acbb000c3d9"
 dependencies = [
  "comfy-table",
  "re_arrow2",
@@ -5891,9 +5738,9 @@ dependencies = [
 
 [[package]]
 name = "re_int_histogram"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bfd863353fcb744faf308e85a70a2e88c539bbaff4f2d32918ec495b5667bba0"
+checksum = "713ac89c0e8143a32f54d13cf9b9bf55a0ad165e8193a1157d920dff29e81e4c"
 dependencies = [
  "smallvec",
  "static_assertions",
@@ -5901,9 +5748,9 @@ dependencies = [
 
 [[package]]
 name = "re_log"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "249e38adcb1a2c69e7a09dbdcced3370c36583471d81773cdedd74df038493b4"
+checksum = "0fc6e1d8b0f15e887dde00051ec776cb64bf03151937aa26e64cfe3c0446e9da"
 dependencies = [
  "env_logger",
  "js-sys",
@@ -5916,9 +5763,9 @@ dependencies = [
 
 [[package]]
 name = "re_log_encoding"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b12929caca153106cc41904cf3a74583b070f628a05ec888f59ce25c74609513"
+checksum = "5b64633b75e8acc2c503311c94fbaccabe217bdd24a811d8e6bfb21f2e1e4a1e"
 dependencies = [
  "lz4_flex",
  "parking_lot",
@@ -5934,9 +5781,9 @@ dependencies = [
 
 [[package]]
 name = "re_log_types"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "99280e77603a8a7b04e2f6149a6773d6ef50079637d763d131839ad78ce16932"
+checksum = "4eae0402f95562af7010696377413205c97fb45bce9de27c03d4678a47be9cdd"
 dependencies = [
  "ahash",
  "anyhow",
@@ -5945,7 +5792,7 @@ dependencies = [
  "clean-path",
  "crossbeam",
  "document-features",
- "fixed 0.5.7",
+ "fixed",
  "half",
  "itertools 0.13.0",
  "natord",
@@ -5973,13 +5820,13 @@ dependencies = [
 
 [[package]]
 name = "re_memory"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bc0ea7c80b0dc8f42598086d926dd6b749c14c4e53ee975e4915e4022caf2cca"
+checksum = "955c483e350ed6d438d64fc3c34081d83b67f5671086a65d46b399e4b30a275a"
 dependencies = [
  "ahash",
  "backtrace",
- "emath 0.29.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "emath",
  "itertools 0.13.0",
  "memory-stats",
  "nohash-hasher",
@@ -5996,13 +5843,12 @@ dependencies = [
 
 [[package]]
 name = "re_mp4"
-version = "0.1.0"
+version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3d1e30657b1ae7f0dd3428a59dc8140732b74a22cc07763606c9ec4054138731"
+checksum = "751650322999417b64a5a89b034b4e34e4596826e5dfee9327856db77ca511e3"
 dependencies = [
  "byteorder",
  "bytes",
- "log",
  "num-rational",
  "serde",
  "serde_json",
@@ -6011,9 +5857,9 @@ dependencies = [
 
 [[package]]
 name = "re_query"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "88d428f95fc875a717186fcf84d3e59f99c4ff76dce7100a89bc440b9ffc7fa9"
+checksum = "6b54efe4a46be8e16921104814b3d91d582d3901e76e127d2dc3d13dc78220c0"
 dependencies = [
  "ahash",
  "anyhow",
@@ -6038,9 +5884,9 @@ dependencies = [
 
 [[package]]
 name = "re_sdk"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "da99da0c9129033de969e91f3ce1d8d6b98ef53321e3760a7fc66aab00cc6a4b"
+checksum = "554227ef7c5ae33279f8e54a004b8d647886ec2d56cc85ce7b353b812933d508"
 dependencies = [
  "ahash",
  "crossbeam",
@@ -6064,9 +5910,9 @@ dependencies = [
 
 [[package]]
 name = "re_sdk_comms"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "daced7adbfd36d957d53adbd91d365715cdb361f071d7f87e09f7df8945e926c"
+checksum = "94cf00b1de4eda76a4e82d82093f5dd44ac210458420ec6a108ad1244b554f82"
 dependencies = [
  "ahash",
  "crossbeam",
@@ -6081,9 +5927,9 @@ dependencies = [
 
 [[package]]
 name = "re_smart_channel"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9f81a0042da6ca0acd4da2c1cbb24c3a91ed426844bacb661a86db1fa9f507d8"
+checksum = "e1508722dacb39ee80a7cae303683826c1fe0a6f6096a1bd27b26097662c7c6f"
 dependencies = [
  "crossbeam",
  "parking_lot",
@@ -6094,9 +5940,9 @@ dependencies = [
 
 [[package]]
 name = "re_string_interner"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "aae7352fe6b1d0dc69a72880a78b0b6c3d165aeff14a7d91102c7b992e23911c"
+checksum = "6bdec7efb723c8306666baec8a0fe924a2dbdc6ae9e21364c171da046a397ff2"
 dependencies = [
  "ahash",
  "nohash-hasher",
@@ -6108,21 +5954,18 @@ dependencies = [
 
 [[package]]
 name = "re_tracing"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "515173071739470a870b960674da3b01190e90c5fd11363552934163320c57c7"
+checksum = "e223d16abbb2aa466eadbd369f982fc4d2c12b7f49f77971c51768a893ca6aca"
 dependencies = [
  "puffin",
- "puffin_http",
- "re_log",
- "rfd",
 ]
 
 [[package]]
 name = "re_tuid"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "533a55427c77321336c7c7fc58e00d0e975e6a8769f3cd586e4f630ac00cdacd"
+checksum = "a3e817c3f2edffa4f04378d33b8ec9b16a1abce2ba28dd6bf0b71c755fc88f20"
 dependencies = [
  "document-features",
  "getrandom",
@@ -6133,16 +5976,16 @@ dependencies = [
 
 [[package]]
 name = "re_types"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f9808ff5fe6e2d585e51fd772bc7bf2320bfe3a38090fc3d088c1b3297c4fe8c"
+checksum = "6c32c586bcf383345c3507808243f8a7fde95e2abcff8e356605a6a644a68112"
 dependencies = [
  "anyhow",
  "array-init",
  "bytemuck",
  "document-features",
- "ecolor 0.29.1 (registry+https://github.com/rust-lang/crates.io-index)",
- "emath 0.29.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "ecolor",
+ "emath",
  "glam",
  "half",
  "image",
@@ -6170,11 +6013,12 @@ dependencies = [
 
 [[package]]
 name = "re_types_builder"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1acbe0cabc9fb933d273604660df9a264f81f2c90938077bfb089765eaeaf1e1"
+checksum = "22482d11e2650298d917016a6b7e696433735acd7bf4e01f2812aecc0626c32d"
 dependencies = [
  "anyhow",
+ "arrow",
  "camino",
  "clang-format",
  "flatbuffers",
@@ -6199,9 +6043,9 @@ dependencies = [
 
 [[package]]
 name = "re_types_core"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3d7508a5209bddc61f48f2af34e372521a568deb53d1b71faf4d17a2e6c24196"
+checksum = "bc8aef506e030a6fd0a2ddfd55b5e9a8babbfb396b9cb5c321b41648c809adc6"
 dependencies = [
  "anyhow",
  "backtrace",
@@ -6223,19 +6067,26 @@ dependencies = [
 
 [[package]]
 name = "re_video"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "13777ccfc3a7fc873dc20af62ca9ae56ac2d7205f5aceb7eb944199aa9f126e8"
+checksum = "31b996b48fa9501ca3f3912c6792538d36b7af4ffef7b20114b23920c0845b42"
 dependencies = [
+ "bit-vec 0.7.0",
  "cfg_aliases 0.2.1",
  "crossbeam",
  "econtext",
  "itertools 0.13.0",
+ "js-sys",
+ "once_cell",
  "parking_lot",
+ "re_build_info",
+ "re_build_tools",
  "re_log",
  "re_mp4",
  "re_tracing",
  "thiserror 1.0.69",
+ "wasm-bindgen",
+ "web-sys",
 ]
 
 [[package]]
@@ -6381,9 +6232,9 @@ dependencies = [
 
 [[package]]
 name = "rerun"
-version = "0.19.1"
+version = "0.20.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b6ce4d70680b035fe98b076e5e2025ae015dc9ac3a08080dd87d3faa4b4653d6"
+checksum = "ed4ff45b89d955f60249837e14ae326eb166e5a3dc4aec653ba917d6c185f4d3"
 dependencies = [
  "anyhow",
  "document-features",
@@ -6414,11 +6265,8 @@ version = "0.12.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3c9e7b57df6e8472152674607f6cc68aa14a748a3157a857a94f516e11aeacc2"
 dependencies = [
- "ashpd",
- "async-io 1.13.0",
  "block",
  "dispatch",
- "futures-util",
  "glib-sys",
  "gobject-sys",
  "gtk-sys",
@@ -6427,9 +6275,7 @@ dependencies = [
  "objc",
  "objc-foundation",
  "objc_id",
- "pollster",
  "raw-window-handle 0.5.2",
- "urlencoding",
  "wasm-bindgen",
  "wasm-bindgen-futures",
  "web-sys",
@@ -6510,7 +6356,7 @@ checksum = "825ea780781b15345a146be27eaefb05085e337e869bff01b4306a4fd4a9ad5a"
 dependencies = [
  "cfg-if",
  "glob",
- "proc-macro-crate 3.2.0",
+ "proc-macro-crate",
  "proc-macro2",
  "quote",
  "regex",
@@ -6567,20 +6413,6 @@ dependencies = [
  "semver",
 ]
 
-[[package]]
-name = "rustix"
-version = "0.37.27"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fea8ca367a3a01fe35e6943c400addf443c0f57670e6ec51196f71a4b8762dd2"
-dependencies = [
- "bitflags 1.3.2",
- "errno",
- "io-lifetimes",
- "libc",
- "linux-raw-sys 0.3.8",
- "windows-sys 0.48.0",
-]
-
 [[package]]
 name = "rustix"
 version = "0.38.41"
@@ -6590,7 +6422,7 @@ dependencies = [
  "bitflags 2.6.0",
  "errno",
  "libc",
- "linux-raw-sys 0.4.14",
+ "linux-raw-sys",
  "windows-sys 0.52.0",
 ]
 
@@ -6788,17 +6620,6 @@ dependencies = [
  "serde",
 ]
 
-[[package]]
-name = "serde_repr"
-version = "0.1.19"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "6c64451ba24fc7a6a2d60fc75dd9c83c90903b19028d4eff35e88fc1e86564e9"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 2.0.90",
-]
-
 [[package]]
 name = "serde_rusqlite"
 version = "0.36.0"
@@ -6990,7 +6811,7 @@ dependencies = [
  "libc",
  "log",
  "memmap2",
- "rustix 0.38.41",
+ "rustix",
  "thiserror 1.0.69",
  "wayland-backend",
  "wayland-client",
@@ -7022,16 +6843,6 @@ dependencies = [
  "serde",
 ]
 
-[[package]]
-name = "socket2"
-version = "0.4.10"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9f7916fc008ca5542385b89a3d3ce689953c143e9304a9bf8beec1de48994c0d"
-dependencies = [
- "libc",
- "winapi",
-]
-
 [[package]]
 name = "socket2"
 version = "0.5.8"
@@ -7291,9 +7102,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "28cce251fcbc87fac86a866eeb0d6c2d536fc16d06f184bb61aeae11aa4cee0c"
 dependencies = [
  "cfg-if",
- "fastrand 2.2.0",
+ "fastrand",
  "once_cell",
- "rustix 0.38.41",
+ "rustix",
  "windows-sys 0.59.0",
 ]
 
@@ -7312,7 +7123,7 @@ version = "0.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5352447f921fda68cf61b4101566c0bdb5104eff6804d0678e5227580ab6a4e9"
 dependencies = [
- "rustix 0.38.41",
+ "rustix",
  "windows-sys 0.59.0",
 ]
 
@@ -7348,11 +7159,11 @@ dependencies = [
 
 [[package]]
 name = "thiserror"
-version = "2.0.3"
+version = "2.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c006c85c7651b3cf2ada4584faa36773bd07bac24acfb39f3c431b36d7e667aa"
+checksum = "2f49a1853cf82743e3b7950f77e0f4d622ca36cf4317cba00c767838bac8d490"
 dependencies = [
- "thiserror-impl 2.0.3",
+ "thiserror-impl 2.0.4",
 ]
 
 [[package]]
@@ -7368,9 +7179,9 @@ dependencies = [
 
 [[package]]
 name = "thiserror-impl"
-version = "2.0.3"
+version = "2.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f077553d607adc1caf65430528a576c757a71ed73944b66ebb58ef2bbd243568"
+checksum = "8381894bb3efe0c4acac3ded651301ceee58a15d47c2e34885ed1908ad667061"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -7409,9 +7220,9 @@ dependencies = [
 
 [[package]]
 name = "time"
-version = "0.3.36"
+version = "0.3.37"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5dfd88e563464686c916c7e46e623e520ddc6d79fa6641390f2e3fa86e83e885"
+checksum = "35e7868883861bd0e56d9ac6efcaaca0d6d5d82a2a7ec8209ff492c07cf37b21"
 dependencies = [
  "deranged",
  "itoa",
@@ -7433,14 +7244,23 @@ checksum = "ef927ca75afb808a4d64dd374f00a2adf8d0fcff8e7b184af886c3c87ec4a3f3"
 
 [[package]]
 name = "time-macros"
-version = "0.2.18"
+version = "0.2.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3f252a68540fde3a3877aeea552b832b40ab9a69e318efd078774a01ddee1ccf"
+checksum = "2834e6017e3e5e4b9834939793b282bc03b37a3336245fa820e35e233e2a85de"
 dependencies = [
  "num-conv",
  "time-core",
 ]
 
+[[package]]
+name = "tiny-keccak"
+version = "2.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237"
+dependencies = [
+ "crunchy",
+]
+
 [[package]]
 name = "tiny-skia"
 version = "0.11.4"
@@ -7493,16 +7313,16 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
 
 [[package]]
 name = "tokio"
-version = "1.41.1"
+version = "1.42.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "22cfb5bee7a6a52939ca9224d6ac897bb669134078daa8735560897f69de4d33"
+checksum = "5cec9b21b0450273377fc97bd4c33a8acffc8c996c987a7c5b319a0083707551"
 dependencies = [
  "backtrace",
  "bytes",
  "libc",
  "mio",
  "pin-project-lite",
- "socket2 0.5.8",
+ "socket2",
  "tokio-macros",
  "windows-sys 0.52.0",
 ]
@@ -7552,9 +7372,9 @@ dependencies = [
 
 [[package]]
 name = "tokio-util"
-version = "0.7.12"
+version = "0.7.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "61e7c3654c13bcd040d4a03abee2c75b1d14a37b423cf5a813ceae1cc903ec6a"
+checksum = "d7fcaa8d55a2bdd6b83ace262b016eca0d79ee02818c5c1bcdf0305114081078"
 dependencies = [
  "bytes",
  "futures-core",
@@ -7596,7 +7416,7 @@ dependencies = [
  "serde",
  "serde_spanned",
  "toml_datetime",
- "toml_edit 0.22.22",
+ "toml_edit",
 ]
 
 [[package]]
@@ -7608,17 +7428,6 @@ dependencies = [
  "serde",
 ]
 
-[[package]]
-name = "toml_edit"
-version = "0.19.15"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "1b5bb770da30e5cbfde35a2d7b9b8a2c4b8ef89548a7a6aeab5c9a576e3e7421"
-dependencies = [
- "indexmap",
- "toml_datetime",
- "winnow 0.5.40",
-]
-
 [[package]]
 name = "toml_edit"
 version = "0.22.22"
@@ -7629,7 +7438,7 @@ dependencies = [
  "serde",
  "serde_spanned",
  "toml_datetime",
- "winnow 0.6.20",
+ "winnow",
 ]
 
 [[package]]
@@ -7831,17 +7640,6 @@ version = "0.10.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f720def6ce1ee2fc44d40ac9ed6d3a59c361c80a75a7aa8e75bb9baed31cf2ea"
 
-[[package]]
-name = "uds_windows"
-version = "1.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "89daebc3e6fd160ac4aa9fc8b3bf71e1f74fbf92367ae71fb83a037e8bf164b9"
-dependencies = [
- "memoffset 0.9.1",
- "tempfile",
- "winapi",
-]
-
 [[package]]
 name = "ug"
 version = "0.0.2"
@@ -7931,9 +7729,9 @@ checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1"
 
 [[package]]
 name = "ureq"
-version = "2.10.1"
+version = "2.12.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b74fc6b57825be3373f7054754755f03ac3a8f5d70015ccad699ba2029956f4a"
+checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d"
 dependencies = [
  "base64",
  "flate2",
@@ -7956,7 +7754,6 @@ dependencies = [
  "form_urlencoded",
  "idna",
  "percent-encoding",
- "serde",
 ]
 
 [[package]]
@@ -8024,12 +7821,6 @@ version = "0.9.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
 
-[[package]]
-name = "waker-fn"
-version = "1.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "317211a0dc0ceedd78fb2ca9a44aed3d7b9b26f81870d485c07122b4350673b7"
-
 [[package]]
 name = "walkdir"
 version = "2.5.0"
@@ -8155,7 +7946,7 @@ checksum = "056535ced7a150d45159d3a8dc30f91a2e2d588ca0b23f70e56033622b8016f6"
 dependencies = [
  "cc",
  "downcast-rs",
- "rustix 0.38.41",
+ "rustix",
  "scoped-tls",
  "smallvec",
  "wayland-sys",
@@ -8168,7 +7959,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "b66249d3fc69f76fd74c82cc319300faa554e9d865dab1f7cd66cc20db10b280"
 dependencies = [
  "bitflags 2.6.0",
- "rustix 0.38.41",
+ "rustix",
  "wayland-backend",
  "wayland-scanner",
 ]
@@ -8190,7 +7981,7 @@ version = "0.31.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "32b08bc3aafdb0035e7fe0fdf17ba0c09c268732707dca4ae098f60cb28c9e4c"
 dependencies = [
- "rustix 0.38.41",
+ "rustix",
  "wayland-client",
  "xcursor",
 ]
@@ -8841,7 +8632,7 @@ dependencies = [
  "pin-project",
  "raw-window-handle 0.6.2",
  "redox_syscall 0.4.1",
- "rustix 0.38.41",
+ "rustix",
  "sctk-adwaita",
  "smithay-client-toolkit",
  "smol_str",
@@ -8861,15 +8652,6 @@ dependencies = [
  "xkbcommon-dl",
 ]
 
-[[package]]
-name = "winnow"
-version = "0.5.40"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "f593a95398737aeed53e489c785df13f3618e41dbcd6718c6addbf1395aa6876"
-dependencies = [
- "memchr",
-]
-
 [[package]]
 name = "winnow"
 version = "0.6.20"
@@ -8925,7 +8707,7 @@ dependencies = [
  "libc",
  "libloading",
  "once_cell",
- "rustix 0.38.41",
+ "rustix",
  "x11rb-protocol",
 ]
 
@@ -8941,16 +8723,6 @@ version = "0.3.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0ef33da6b1660b4ddbfb3aef0ade110c8b8a781a3b6382fa5f2b5b040fd55f61"
 
-[[package]]
-name = "xdg-home"
-version = "1.3.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ec1cdab258fb55c0da61328dc52c8764709b249011b2cad0454c72f0bf10a1f6"
-dependencies = [
- "libc",
- "windows-sys 0.59.0",
-]
-
 [[package]]
 name = "xkbcommon-dl"
 version = "0.4.2"
@@ -8972,9 +8744,9 @@ checksum = "b9cc00251562a284751c9973bace760d86c0276c471b4be569fe6b068ee97a56"
 
 [[package]]
 name = "xml-rs"
-version = "0.8.23"
+version = "0.8.24"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "af310deaae937e48a26602b730250b4949e125f468f11e6990be3e5304ddd96f"
+checksum = "ea8b391c9a790b496184c29f7f93b9ed5b16abb306c05415b68bcc16e4d06432"
 
 [[package]]
 name = "xshell"
@@ -9021,72 +8793,6 @@ dependencies = [
  "synstructure",
 ]
 
-[[package]]
-name = "zbus"
-version = "3.15.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "675d170b632a6ad49804c8cf2105d7c31eddd3312555cffd4b740e08e97c25e6"
-dependencies = [
- "async-broadcast",
- "async-executor",
- "async-fs 1.6.0",
- "async-io 1.13.0",
- "async-lock 2.8.0",
- "async-process",
- "async-recursion",
- "async-task",
- "async-trait",
- "blocking",
- "byteorder",
- "derivative",
- "enumflags2",
- "event-listener 2.5.3",
- "futures-core",
- "futures-sink",
- "futures-util",
- "hex",
- "nix",
- "once_cell",
- "ordered-stream",
- "rand",
- "serde",
- "serde_repr",
- "sha1",
- "static_assertions",
- "tracing",
- "uds_windows",
- "winapi",
- "xdg-home",
- "zbus_macros",
- "zbus_names",
- "zvariant",
-]
-
-[[package]]
-name = "zbus_macros"
-version = "3.15.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7131497b0f887e8061b430c530240063d33bf9455fa34438f388a245da69e0a5"
-dependencies = [
- "proc-macro-crate 1.3.1",
- "proc-macro2",
- "quote",
- "regex",
- "syn 1.0.109",
- "zvariant_utils",
-]
-
-[[package]]
-name = "zbus_names"
-version = "2.6.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "437d738d3750bed6ca9b8d423ccc7a8eb284f6b1d6d4e225a0e4e6258d864c8d"
-dependencies = [
- "serde",
- "static_assertions",
- "zvariant",
-]
-
 [[package]]
 name = "zerocopy"
 version = "0.7.35"
@@ -9205,7 +8911,7 @@ dependencies = [
  "flate2",
  "indexmap",
  "memchr",
- "thiserror 2.0.3",
+ "thiserror 2.0.4",
  "zopfli",
 ]
 
@@ -9275,42 +8981,3 @@ checksum = "16099418600b4d8f028622f73ff6e3deaabdff330fb9a2a131dea781ee8b0768"
 dependencies = [
  "zune-core",
 ]
-
-[[package]]
-name = "zvariant"
-version = "3.15.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4eef2be88ba09b358d3b58aca6e41cd853631d44787f319a1383ca83424fb2db"
-dependencies = [
- "byteorder",
- "enumflags2",
- "libc",
- "serde",
- "static_assertions",
- "url",
- "zvariant_derive",
-]
-
-[[package]]
-name = "zvariant_derive"
-version = "3.15.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "37c24dc0bed72f5f90d1f8bb5b07228cbf63b3c6e9f82d82559d4bae666e7ed9"
-dependencies = [
- "proc-macro-crate 1.3.1",
- "proc-macro2",
- "quote",
- "syn 1.0.109",
- "zvariant_utils",
-]
-
-[[package]]
-name = "zvariant_utils"
-version = "1.0.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7234f0d811589db492d16893e3f21e8e2fd282e6d01b0cddee310322062cc200"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 1.0.109",
-]
diff --git a/Cargo.toml b/Cargo.toml
index af8e9a8b..80a118f3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -27,7 +27,7 @@ serde_json = { version = "1.0.128", default-features = false }
 
 rand = "0.8.5"
 anyhow = "1.0.81"
-tracing = "0.1.40"
+tracing = "0.1.30"
 tracing-tracy = "0.11.0"
 tracing-subscriber = "0.3.18"
 
@@ -78,6 +78,7 @@ burn-wgpu = { git = "https://github.com/tracel-ai/burn", features = [
 burn-fusion = { git = "https://github.com/tracel-ai/burn" }
 
 egui = { git = "https://github.com/emilk/egui/", rev = "5bfff316c9818b3c140d02bb6cdc488556d46ab7" }
+
 eframe = { git = "https://github.com/emilk/egui/", rev = "5bfff316c9818b3c140d02bb6cdc488556d46ab7", default-features = false, features = [
     "wgpu",
     "android-game-activity",
@@ -89,7 +90,7 @@ egui_extras = { git = "https://github.com/emilk/egui/", rev = "5bfff316c9818b3c1
     "all_loaders",
 ] }
 
-rerun = { version = "0.19.1", default-features = false, features = [
+rerun = { version = "0.20.0", default-features = false, features = [
     'sdk',
     'glam',
     'image',
@@ -107,6 +108,7 @@ web-sys = { version = "0.3.72", features = [
 wasm-logger = "0.2.0"
 zip = { version = "2.1.3", default-features = false, features = ["deflate"] }
 urlencoding = "2.1"
+hashbrown = "0.15"
 
 [patch."https://github.com/tracel-ai/burn"]
 # Uncomment this to use local burn.
@@ -124,6 +126,9 @@ urlencoding = "2.1"
 wgpu = { git = "https://github.com/ArthurBrussee/wgpu", branch = "flt-atom" }
 naga = { git = "https://github.com/ArthurBrussee/wgpu", branch = "flt-atom" }
 
+emath = { git = "https://github.com/emilk/egui/", rev = "5bfff316c9818b3c140d02bb6cdc488556d46ab7" }
+ecolor = { git = "https://github.com/emilk/egui/", rev = "5bfff316c9818b3c140d02bb6cdc488556d46ab7" }
+
 [profile.dev]
 opt-level = 0
 debug = true
diff --git a/crates/brush-dataset/src/lib.rs b/crates/brush-dataset/src/lib.rs
index 5a182ee6..c0661d36 100644
--- a/crates/brush-dataset/src/lib.rs
+++ b/crates/brush-dataset/src/lib.rs
@@ -33,7 +33,7 @@ pub struct LoadInitArgs {
 
 impl Default for LoadInitArgs {
     fn default() -> Self {
-        Self { sh_degree: 2 }
+        Self { sh_degree: 3 }
     }
 }
 
@@ -90,7 +90,7 @@ pub(crate) fn stream_fut_parallel<T: Send + 'static>(
             .get()
     };
 
-    log::info!("Loading steam with {parallel} threads");
+    log::info!("Loading stream with {parallel} threads");
 
     let mut futures = futures;
     fn_stream(|emitter| async move {
diff --git a/crates/brush-dataset/src/scene_loader.rs b/crates/brush-dataset/src/scene_loader.rs
index 0b5be796..e23e8170 100644
--- a/crates/brush-dataset/src/scene_loader.rs
+++ b/crates/brush-dataset/src/scene_loader.rs
@@ -18,7 +18,16 @@ impl<B: Backend> SceneLoader<B> {
         // The bounded size == number of batches to prefetch.
         let (tx, rx) = mpsc::channel(5);
         let device = device.clone();
-        let scene_extent = scene.bounds(0.0, 0.0).extent.max_element() as f64;
+
+        let center = scene.bounds().center;
+        let dists = scene
+            .views
+            .iter()
+            .map(|v| (v.camera.position - center).length())
+            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Less))
+            .unwrap();
+
+        let scene_extent = dists * 1.1; // Idk why exactly, but gsplat multiplies this by 1.1
 
         let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
 
diff --git a/crates/brush-render/src/burn_glue.rs b/crates/brush-render/src/burn_glue.rs
index 866af885..321113f6 100644
--- a/crates/brush-render/src/burn_glue.rs
+++ b/crates/brush-render/src/burn_glue.rs
@@ -8,8 +8,9 @@ use burn::{
         Autodiff,
     },
     tensor::{
+        backend::AutodiffBackend,
         repr::{CustomOpDescription, HandleContainer, OperationDescription},
-        BasicAutodiffOps, DType, Float, Tensor, TensorPrimitive,
+        DType, Tensor, TensorPrimitive,
     },
 };
 use burn_fusion::{client::FusionClient, stream::Operation, Fusion};
@@ -22,7 +23,7 @@ use crate::{
         calc_tile_bounds, max_intersections, render_backward, render_forward, sh_coeffs_for_degree,
         sh_degree_from_coeffs,
     },
-    shaders, AutodiffBackend, Backend, GaussianBackwardState, InnerWgpu, RenderAux, SplatGrads,
+    shaders, Backend, GaussianBackwardState, InnerWgpu, RenderAux, SplatGrads,
 };
 
 // Implement forward functions for the inner wgpu backend.
@@ -170,15 +171,10 @@ impl<B: Backend, C: CheckpointStrategy> Backend for Autodiff<B, C> {
             render_u32_buffer,
         );
 
-        // Not sure why going into the autodiff float tensor type is so verbose.
-        let diff_proj = <Float as BasicAutodiffOps<Self>>::from_inner(TensorPrimitive::Float(
-            aux.projected_splats.clone(),
-        ))
-        .tensor();
-
         let auxc = aux.clone();
         let wrapped_aux = RenderAux::<Self> {
-            projected_splats: diff_proj,
+            projected_splats: <Self as AutodiffBackend>::from_inner(aux.projected_splats),
+            radii: <Self as AutodiffBackend>::from_inner(aux.radii),
             num_intersections: aux.num_intersections,
             num_visible: aux.num_visible,
             final_index: aux.final_index,
@@ -241,7 +237,7 @@ impl Backend for Fusion<InnerWgpu> {
             fn execute(self: Box<Self>, h: &mut HandleContainer<JitFusionHandle<WgpuRuntime>>) {
                 let (
                     [means, log_scales, quats, sh_coeffs, raw_opacity],
-                    [projected_splats, uniforms_buffer, num_intersections, num_visible, final_index, cum_tiles_hit, tile_bins, compact_gid_from_isect, global_from_compact_gid, out_img],
+                    [projected_splats, uniforms_buffer, num_intersections, num_visible, final_index, cum_tiles_hit, tile_bins, compact_gid_from_isect, global_from_compact_gid, radii, out_img],
                 ) = self.desc.consume();
 
                 let (img, aux) = render_forward(
@@ -272,6 +268,7 @@ impl Backend for Fusion<InnerWgpu> {
                     &global_from_compact_gid.id,
                     aux.global_from_compact_gid,
                 );
+                h.register_float_tensor::<InnerWgpu>(&radii.id, aux.radii);
             }
         }
 
@@ -309,6 +306,7 @@ impl Backend for Fusion<InnerWgpu> {
             compact_gid_from_isect: client
                 .tensor_uninitialized(vec![max_intersects as usize], DType::I32),
             global_from_compact_gid: client.tensor_uninitialized(vec![num_points], DType::I32),
+            radii: client.tensor_uninitialized(vec![num_points], DType::F32),
         };
 
         let desc = CustomOpDescription::new(
@@ -330,6 +328,7 @@ impl Backend for Fusion<InnerWgpu> {
                 aux.tile_bins.to_description_out(),
                 aux.compact_gid_from_isect.to_description_out(),
                 aux.global_from_compact_gid.to_description_out(),
+                aux.radii.to_description_out(),
                 out_img.to_description_out(),
             ],
         );
@@ -441,4 +440,4 @@ impl Backend for Fusion<InnerWgpu> {
     }
 }
 
-impl<B: Backend, C: CheckpointStrategy> AutodiffBackend for Autodiff<B, C> {}
+impl<B: Backend, C: CheckpointStrategy> crate::AutodiffBackend for Autodiff<B, C> {}
diff --git a/crates/brush-render/src/gaussian_splats.rs b/crates/brush-render/src/gaussian_splats.rs
index 8082c70f..7fee936e 100644
--- a/crates/brush-render/src/gaussian_splats.rs
+++ b/crates/brush-render/src/gaussian_splats.rs
@@ -107,20 +107,21 @@ impl<B: Backend> Splats<B> {
                 .iter()
                 .map(|p| {
                     // Get average of 3 nearest squared distances.
-                    tree.nearest_n::<SquaredEuclidean>(p, 3)
+                    (tree
+                        .nearest_n::<SquaredEuclidean>(p, 4)
                         .iter()
                         .map(|x| x.distance)
                         .sum::<f32>()
+                        / 4.0)
                         .sqrt()
-                        / 3.0
+                        .max(1e-12)
+                        .ln()
                 })
                 .collect();
 
             Tensor::<B, 1>::from_floats(extents.as_slice(), device)
                 .reshape([n_splats, 1])
                 .repeat_dim(1, 3)
-                .clamp_min(0.00001)
-                .log()
         };
 
         let sh_coeffs = if let Some(sh_coeffs) = sh_coeffs {
@@ -150,7 +151,7 @@ impl<B: Backend> Splats<B> {
         )
     }
 
-    pub fn with_min_sh_degree(mut self, sh_degree: u32) -> Self {
+    pub fn with_sh_degree(mut self, sh_degree: u32) -> Self {
         let n_coeffs = sh_coeffs_for_degree(sh_degree) as usize;
 
         let [n, c, _] = self.sh_coeffs.dims();
diff --git a/crates/brush-render/src/lib.rs b/crates/brush-render/src/lib.rs
index 68268e62..c3151058 100644
--- a/crates/brush-render/src/lib.rs
+++ b/crates/brush-render/src/lib.rs
@@ -29,6 +29,7 @@ pub struct RenderAux<B: Backend> {
     pub tile_bins: B::IntTensorPrimitive,
     pub compact_gid_from_isect: B::IntTensorPrimitive,
     pub global_from_compact_gid: B::IntTensorPrimitive,
+    pub radii: B::FloatTensorPrimitive,
 }
 
 #[derive(Debug, Clone)]
@@ -115,5 +116,8 @@ pub trait Backend: burn::tensor::backend::Backend {
     }
 }
 
-pub trait AutodiffBackend: Backend + burn::tensor::backend::AutodiffBackend {}
+pub trait AutodiffBackend:
+    Backend + burn::tensor::backend::AutodiffBackend<InnerBackend: Backend>
+{
+}
 type InnerWgpu = JitBackend<WgpuRuntime, f32, i32, u32>;
diff --git a/crates/brush-render/src/render.rs b/crates/brush-render/src/render.rs
index 41190f80..72bb1c2d 100644
--- a/crates/brush-render/src/render.rs
+++ b/crates/brush-render/src/render.rs
@@ -154,6 +154,8 @@ pub(crate) fn render_forward(
         create_tensor::<1, WgpuRuntime>([num_points], device, client, DType::U32);
     let num_tiles = InnerWgpu::int_zeros([num_points].into(), device);
 
+    let radii = InnerWgpu::float_zeros([num_points].into(), device);
+
     let (global_from_compact_gid, num_visible) = {
         let global_from_presort_gid = InnerWgpu::int_zeros([num_points].into(), device);
         let depths = create_tensor([num_points], device, client, DType::F32);
@@ -172,7 +174,8 @@ pub(crate) fn render_forward(
                     raw_opacities.clone().handle.binding(),
                     global_from_presort_gid.clone().handle.binding(),
                     depths.clone().handle.binding(),
-                    num_tiles_scatter.clone().handle.binding()
+                    num_tiles_scatter.clone().handle.binding(),
+                    radii.clone().handle.binding(),
                 ],
             );
         });
@@ -354,6 +357,7 @@ pub(crate) fn render_forward(
             final_index,
             compact_gid_from_isect,
             global_from_compact_gid,
+            radii,
         },
     )
 }
@@ -384,7 +388,7 @@ pub(crate) fn render_backward(
 
     let client = &means.client;
 
-    let (v_xys_local, v_xys_global, v_conics, v_coeffs, v_raw_opac) = {
+    let (v_xys_local, v_conics, v_coeffs, v_raw_opac) = {
         let tile_bounds = uvec2(
             img_size.x.div_ceil(shaders::helpers::TILE_WIDTH),
             img_size.y.div_ceil(shaders::helpers::TILE_WIDTH),
@@ -430,7 +434,6 @@ pub(crate) fn render_backward(
 
         let num_vis_wg = create_dispatch_buffer(num_visible.clone(), GatherGrads::WORKGROUP_SIZE);
 
-        let v_xys_global = InnerWgpu::float_zeros([num_points, 2].into(), device);
         unsafe {
             client.execute_unchecked(
                 GatherGrads::task(),
@@ -441,15 +444,13 @@ pub(crate) fn render_backward(
                     raw_opac.clone().handle.binding(),
                     means.clone().handle.binding(),
                     v_colors.clone().handle.binding(),
-                    v_xys_local.clone().handle.binding(),
                     v_coeffs.handle.clone().binding(),
                     v_opacities.handle.clone().binding(),
-                    v_xys_global.handle.clone().binding(),
                 ],
             );
         }
 
-        (v_xys_local, v_xys_global, v_conics, v_coeffs, v_opacities)
+        (v_xys_local, v_conics, v_coeffs, v_opacities)
     };
 
     // Create tensors to hold gradients.
@@ -485,7 +486,7 @@ pub(crate) fn render_backward(
         v_scales,
         v_coeffs,
         v_raw_opac,
-        v_xy: v_xys_global,
+        v_xy: v_xys_local,
     }
 }
 
@@ -592,7 +593,7 @@ mod tests {
 
         let rec = if USE_RERUN {
             rerun::RecordingStreamBuilder::new("render test")
-                .connect()
+                .connect_tcp()
                 .ok()
         } else {
             None
@@ -688,9 +689,16 @@ mod tests {
                 .mean()
                 .backward();
 
+            // XY gradients are also in compact format.
+            let v_xys = splats.xys_dummy.grad(&grads).context("no xys grad")?;
+            let v_xys = v_xys.slice([0..num_visible]);
+
             let v_xys_ref =
                 safetensor_to_burn::<DiffBack, 2>(tensors.tensor("v_xy")?, &device).inner();
-            let v_xys = splats.xys_dummy.grad(&grads).context("no xys grad")?;
+            let v_xys_ref = v_xys_ref
+                .select(0, gs_ids.inner().clone())
+                .slice([0..num_visible]);
+
             assert!(v_xys.all_close(v_xys_ref, Some(1e-5), Some(1e-9)));
 
             let v_opacities_ref =
diff --git a/crates/brush-render/src/shaders/gather_grads.wgsl b/crates/brush-render/src/shaders/gather_grads.wgsl
index 9e9445e9..30a4f32b 100644
--- a/crates/brush-render/src/shaders/gather_grads.wgsl
+++ b/crates/brush-render/src/shaders/gather_grads.wgsl
@@ -6,11 +6,9 @@
 @group(0) @binding(2) var<storage, read> raw_opacities: array<f32>;
 @group(0) @binding(3) var<storage, read> means: array<helpers::PackedVec3>;
 @group(0) @binding(4) var<storage, read> v_colors: array<vec4f>;
-@group(0) @binding(5) var<storage, read> v_xy_local: array<vec2f>;
 
-@group(0) @binding(6) var<storage, read_write> v_coeffs: array<f32>;
-@group(0) @binding(7) var<storage, read_write> v_opacs: array<f32>;
-@group(0) @binding(8) var<storage, read_write> v_xy_global: array<vec2f>;
+@group(0) @binding(5) var<storage, read_write> v_coeffs: array<f32>;
+@group(0) @binding(6) var<storage, read_write> v_opacs: array<f32>;
 
 const SH_C0: f32 = 0.2820947917738781f;
 
@@ -223,8 +221,4 @@ fn main(@builtin(global_invocation_id) gid: vec3u) {
     let raw_opac = raw_opacities[global_gid];
     let v_opac = v_color.w * v_sigmoid(raw_opac);
     v_opacs[global_gid] = v_opac;
-
-    // Scatter the xy gradients, as later operations need them to be global.
-    let v_xy_local = v_xy_local[compact_gid];
-    v_xy_global[global_gid] = v_xy_local;
 }
diff --git a/crates/brush-render/src/shaders/project_forward.wgsl b/crates/brush-render/src/shaders/project_forward.wgsl
index b5e44e99..b26a424b 100644
--- a/crates/brush-render/src/shaders/project_forward.wgsl
+++ b/crates/brush-render/src/shaders/project_forward.wgsl
@@ -14,6 +14,8 @@
 @group(0) @binding(6) var<storage, read_write> depths: array<f32>;
 @group(0) @binding(7) var<storage, read_write> num_tiles: array<u32>;
 
+@group(0) @binding(8) var<storage, read_write> radii: array<f32>;
+
 @compute
 @workgroup_size(256, 1, 1)
 fn main(@builtin(global_invocation_id) global_id: vec3u) {
@@ -31,7 +33,7 @@ fn main(@builtin(global_invocation_id) global_id: vec3u) {
     let R = mat3x3f(viewmat[0].xyz, viewmat[1].xyz, viewmat[2].xyz);
     let mean_c = R * mean + viewmat[3].xyz;
 
-    if mean_c.z < 0.01 || mean_c.z > 1e12 {
+    if mean_c.z < 0.01 || mean_c.z > 1e10 {
         return;
     }
 
@@ -52,8 +54,6 @@ fn main(@builtin(global_invocation_id) global_id: vec3u) {
 
     // compute the projected mean
     let mean2d = uniforms.focal * mean_c.xy * (1.0 / mean_c.z) + uniforms.pixel_center;
-
-
     let opac = helpers::sigmoid(raw_opacities[global_gid]);
 
     // NB: It might seem silly to use the inverse of the conic here (as that's the same as cov2d)
@@ -91,4 +91,6 @@ fn main(@builtin(global_invocation_id) global_id: vec3u) {
     depths[write_id] = mean_c.z;
     // Write metadata to global array.
     num_tiles[global_gid] = tile_area;
+
+    radii[global_gid] = radius;
 }
diff --git a/crates/brush-train/Cargo.toml b/crates/brush-train/Cargo.toml
index 38697e8e..71364888 100644
--- a/crates/brush-train/Cargo.toml
+++ b/crates/brush-train/Cargo.toml
@@ -14,5 +14,6 @@ glam.workspace = true
 rand.workspace = true
 tracing.workspace = true
 log.workspace = true
+hashbrown.workspace = true
 
 burn.workspace = true
diff --git a/crates/brush-train/src/scene.rs b/crates/brush-train/src/scene.rs
index 7793dbab..8c4fd71f 100644
--- a/crates/brush-train/src/scene.rs
+++ b/crates/brush-train/src/scene.rs
@@ -39,7 +39,13 @@ impl Scene {
     }
 
     // Returns the extent of the cameras in the scene.
-    pub fn bounds(&self, cam_near: f32, cam_far: f32) -> BoundingBox {
+    pub fn bounds(&self) -> BoundingBox {
+        self.adjusted_bounds(0.0, 0.0)
+    }
+
+    // Returns the extent of the cameras in the scene, taking into account
+    // the near and far plane of the cameras.
+    pub fn adjusted_bounds(&self, cam_near: f32, cam_far: f32) -> BoundingBox {
         let (min, max) = self.views.iter().fold(
             (Vec3::splat(f32::INFINITY), Vec3::splat(f32::NEG_INFINITY)),
             |(min, max), view| {
diff --git a/crates/brush-train/src/train.rs b/crates/brush-train/src/train.rs
index 572dd07b..70e939aa 100644
--- a/crates/brush-train/src/train.rs
+++ b/crates/brush-train/src/train.rs
@@ -3,86 +3,86 @@ use brush_render::gaussian_splats::{inverse_sigmoid, Splats};
 use brush_render::{AutodiffBackend, Backend, RenderAux};
 use burn::lr_scheduler::exponential::{ExponentialLrScheduler, ExponentialLrSchedulerConfig};
 use burn::lr_scheduler::LrScheduler;
+use burn::module::{Param, ParamId};
 use burn::optim::adaptor::OptimizerAdaptor;
-use burn::optim::Adam;
-use burn::tensor::{Bool, Distribution, Int};
+use burn::optim::record::AdaptorRecord;
+use burn::optim::{Adam, AdamState};
+use burn::tensor::{Bool, Distribution, Int, TensorPrimitive};
 use burn::{
     config::Config,
     optim::{AdamConfig, GradientsParams, Optimizer},
     tensor::Tensor,
 };
+use hashbrown::HashMap;
 use tracing::trace_span;
 
 use crate::scene::SceneView;
 use crate::ssim::Ssim;
 
+type OptimizerType<B> = OptimizerAdaptor<Adam, Splats<B>, B>;
+
 #[derive(Config)]
 pub struct TrainConfig {
-    // period of steps where refinement is turned off
-    #[config(default = 500)]
-    warmup_steps: u32,
+    // Weight for SSIM loss
+    #[config(default = 0.2)]
+    ssim_weight: f32,
 
-    // period of steps where gaussians are culled and densified
-    #[config(default = 100)]
-    refine_every: u32,
+    // GSs with opacity below this value will be pruned
+    #[config(default = 0.005)]
+    cull_opacity: f32,
 
-    #[config(default = 15000)]
-    max_refine_step: u32,
+    // threshold of positional gradient norm for densifying gaussians
+    #[config(default = 0.0002)]
+    densify_grad_thresh: f32,
 
-    #[config(default = 0.006)]
-    reset_alpha_value: f32,
+    // Gaussians bigger than this size in screenspace radius are split.
+    // Set to 1.0 to disable.
+    #[config(default = 1.0)]
+    densify_radius_threshold: f32,
 
-    // threshold of opacity for culling gaussians. One can set it to a lower value (e.g. 0.005) for higher quality
-    #[config(default = 0.005)]
-    cull_alpha_thresh: f32,
+    // below this size, gaussians are *duplicated*, otherwise split.
+    #[config(default = 0.01)]
+    densify_size_threshold: f32,
 
     // threshold of scale for culling huge gaussians
-    #[config(default = 5.0)]
-    cull_scale_thresh: f32,
+    #[config(default = 0.5)]
+    cull_scale3d_percentage_threshold: f32,
+
+    // period of steps where refinement is turned off
+    #[config(default = 500)]
+    refine_start_iter: u32,
+
+    #[config(default = 15000)]
+    refine_stop_iter: u32,
 
     // Every this many refinement steps, reset the alpha
     #[config(default = 30)]
     reset_alpha_every_refine: u32,
+    // period of steps where gaussians are culled and densified
+    #[config(default = 100)]
+    refine_every: u32,
 
-    // threshold of positional gradient norm for densifying gaussians
-    // TODO: Abs grad.
-    #[config(default = 0.0002)]
-    densify_grad_thresh: f32,
-
-    // below this size, gaussians are *duplicated*, otherwise split.
-    #[config(default = 0.005)]
-    densify_size_thresh: f32,
-
-    #[config(default = 0.1)]
-    ssim_weight: f32,
-
-    #[config(default = 7)]
+    #[config(default = 11)]
     ssim_window_size: usize,
 
-    #[config(default = true)]
-    scale_mean_lr_by_extent: bool,
-
     // Learning rates.
     lr_mean: ExponentialLrSchedulerConfig,
 
-    #[config(default = 0.9999)]
-    lr_global_decay: f64,
-
     // Learning rate for the basic coefficients.
-    #[config(default = 0.006)]
+    #[config(default = 2.5e-3)]
     lr_coeffs_dc: f64,
 
     // How much to divide the learning rate by for higher SH orders.
-    #[config(default = 15.0)]
+    #[config(default = 20.0)]
     lr_coeffs_sh_scale: f64,
 
-    #[config(default = 0.1)]
+    #[config(default = 5e-2)]
     lr_opac: f64,
 
-    #[config(default = 0.02)]
+    #[config(default = 5e-3)]
     lr_scale: f64,
 
-    #[config(default = 0.01)]
+    #[config(default = 1e-3)]
     lr_rotation: f64,
 
     #[config(default = 42)]
@@ -92,8 +92,8 @@ pub struct TrainConfig {
 impl Default for TrainConfig {
     fn default() -> Self {
         let decay_steps = 30000;
-        let lr_max = 3e-4;
-        let decay = 1e-1f64.powf(1.0 / decay_steps as f64);
+        let lr_max = 1.6e-4;
+        let decay = 1e-2f64.powf(1.0 / decay_steps as f64);
         TrainConfig::new(ExponentialLrSchedulerConfig::new(lr_max, decay))
     }
 }
@@ -102,7 +102,7 @@ impl Default for TrainConfig {
 pub struct SceneBatch<B: Backend> {
     pub gt_images: Tensor<B, 4>,
     pub gt_views: Vec<SceneView>,
-    pub scene_extent: f64,
+    pub scene_extent: f32,
 }
 
 #[derive(Clone)]
@@ -129,22 +129,19 @@ pub struct TrainStepStats<B: AutodiffBackend> {
     pub refine: Option<RefineStats>,
 }
 
-pub struct SplatTrainer<B: AutodiffBackend>
-where
-    B::InnerBackend: Backend,
-{
+pub struct SplatTrainer<B: AutodiffBackend> {
     pub iter: u32,
 
     config: TrainConfig,
 
     sched_mean: ExponentialLrScheduler,
-    optim: OptimizerAdaptor<Adam, Splats<B>, B>,
-    opt_config: AdamConfig,
+    optim: OptimizerType<B>,
 
     // Helper tensors for accumulating the viewspace_xy gradients and the number
     // of observations per gaussian. Used in pruning and densification.
     grad_2d_accum: Tensor<B, 1>,
     xy_grad_counts: Tensor<B, 1, Int>,
+    max_radii: Tensor<B, 1>,
 
     ssim: Ssim<B>,
 }
@@ -155,56 +152,60 @@ fn quaternion_vec_multiply<B: Backend>(
 ) -> Tensor<B, 2> {
     let num_points = quaternions.dims()[0];
 
-    // Extract components of quaternions
+    // Extract components
     let qw = quaternions.clone().slice([0..num_points, 0..1]);
     let qx = quaternions.clone().slice([0..num_points, 1..2]);
     let qy = quaternions.clone().slice([0..num_points, 2..3]);
     let qz = quaternions.clone().slice([0..num_points, 3..4]);
 
-    // Extract components of vectors
     let vx = vectors.clone().slice([0..num_points, 0..1]);
     let vy = vectors.clone().slice([0..num_points, 1..2]);
     let vz = vectors.clone().slice([0..num_points, 2..3]);
 
-    // Compute intermediate terms
-    let term1 = qw.clone() * vx.clone() + qy.clone() * vz.clone() - qz.clone() * vy.clone();
-    let term2 = qw.clone() * vy.clone() - qx.clone() * vz.clone() + qz.clone() * vx.clone();
-    let term3 = qw.clone() * vz.clone() + qx.clone() * vy.clone() - qy.clone() * vx.clone();
-    let term4 = qx.clone() * vx.clone() + qy.clone() * vy.clone() + qz.clone() * vz.clone();
-
-    // Compute final result
-    let rx = vx
-        + (qw.clone() * term1.clone() + qx.clone() * term4.clone() - qy.clone() * term3.clone()
-            + qz.clone() * term2.clone())
+    // Common terms
+    let qw2 = qw.clone().powf_scalar(2.0);
+    let qx2 = qx.clone().powf_scalar(2.0);
+    let qy2 = qy.clone().powf_scalar(2.0);
+    let qz2 = qz.clone().powf_scalar(2.0);
+
+    // Cross products (multiplied by 2.0 later)
+    let xy = qx.clone() * qy.clone();
+    let xz = qx.clone() * qz.clone();
+    let yz = qy.clone() * qz.clone();
+    let wx = qw.clone() * qx.clone();
+    let wy = qw.clone() * qy.clone();
+    let wz = qw * qz;
+
+    // Final components with reused terms
+    let x = (qw2.clone() + qx2.clone() - qy2.clone() - qz2.clone()) * vx.clone()
+        + (xy.clone() * vy.clone() + xz.clone() * vz.clone() + wy.clone() * vz.clone()
+            - wz.clone() * vy.clone())
             * 2.0;
-    let ry = vy
-        + (qw.clone() * term2.clone() - qx.clone() * term3.clone()
-            + qy.clone() * term4.clone()
-            + qz.clone() * term1.clone())
+
+    let y = (qw2.clone() - qx2.clone() + qy2.clone() - qz2.clone()) * vy.clone()
+        + (xy.clone() * vx.clone() + yz.clone() * vz.clone() + wz.clone() * vx.clone()
+            - wx.clone() * vz.clone())
             * 2.0;
-    let rz = vz
-        + (qw * term3.clone() + qx * term2.clone() - qy * term1.clone() + qz * term4.clone()) * 2.0;
 
-    Tensor::cat(vec![rx, ry, rz], 1)
+    let z = (qw2 - qx2 - qy2 + qz2) * vz.clone()
+        + (xz * vx.clone() + yz * vy.clone() + wx * vy - wy * vx) * 2.0;
+
+    Tensor::cat(vec![x, y, z], 1)
 }
 
-impl<B: AutodiffBackend> SplatTrainer<B>
-where
-    B::InnerBackend: Backend,
-{
+impl<B: AutodiffBackend> SplatTrainer<B> {
     pub fn new(num_points: usize, config: &TrainConfig, device: &B::Device) -> Self {
         let opt_config = AdamConfig::new().with_epsilon(1e-15);
-        let optim = opt_config.init::<B, Splats<B>>();
-
         let ssim = Ssim::new(config.ssim_window_size, 3, device);
+
         Self {
             config: config.clone(),
             iter: 0,
             sched_mean: config.lr_mean.init().expect("Lr schedule must be valid."),
-            optim,
-            opt_config,
+            optim: opt_config.init::<B, Splats<B>>(),
             grad_2d_accum: Tensor::zeros([num_points], device),
             xy_grad_counts: Tensor::zeros([num_points], device),
+            max_radii: Tensor::zeros([num_points], device),
             ssim,
         }
     }
@@ -212,12 +213,20 @@ where
     fn reset_stats(&mut self, num_points: usize, device: &B::Device) {
         self.grad_2d_accum = Tensor::zeros([num_points], device);
         self.xy_grad_counts = Tensor::zeros([num_points], device);
+        self.max_radii = Tensor::zeros([num_points], device);
     }
 
-    pub(crate) fn reset_opacity(&self, splats: &mut Splats<B>) {
-        Splats::map_param(&mut splats.raw_opacity, |op| {
-            Tensor::zeros_like(&op) + inverse_sigmoid(self.config.reset_alpha_value)
-        });
+    pub(crate) fn reset_opacity(
+        &self,
+        splats: &mut Splats<B>,
+        record: &mut HashMap<ParamId, AdaptorRecord<Adam, B>>,
+    ) {
+        map_param(
+            &mut splats.raw_opacity,
+            record,
+            |op| Tensor::zeros_like(&op) + inverse_sigmoid(self.config.cull_opacity * 2.0),
+            |state| Tensor::zeros_like(&state),
+        );
     }
 
     pub async fn step(
@@ -274,8 +283,8 @@ where
                         .clone()
                         .slice([0..batch_size, 0..img_h, 0..img_w, 0..3]);
 
-                let ssim_loss = self.ssim.ssim(pred_rgb, gt_rgb);
-                loss * (1.0 - self.config.ssim_weight) - ssim_loss * self.config.ssim_weight
+                let ssim_loss = -self.ssim.ssim(pred_rgb, gt_rgb) + 1.0;
+                loss * (1.0 - self.config.ssim_weight) + ssim_loss * self.config.ssim_weight
             } else {
                 loss
             };
@@ -287,26 +296,16 @@ where
 
         // TODO: Should scale lr be scales by scene scale as well?
         let (lr_mean, lr_rotation, lr_scale, lr_coeffs, lr_opac) = (
-            self.sched_mean.step() * batch.scene_extent,
+            self.sched_mean.step() * batch.scene_extent as f64,
             self.config.lr_rotation,
             self.config.lr_scale,
             self.config.lr_coeffs_dc,
             self.config.lr_opac,
         );
 
-        let global_decay = self.config.lr_global_decay.powf(self.iter as f64);
-
-        let (lr_mean, lr_rotation, lr_scale, lr_coeffs, lr_opac) = (
-            lr_mean * global_decay,
-            lr_rotation * global_decay,
-            lr_scale * global_decay,
-            lr_coeffs * global_decay,
-            lr_opac * global_decay,
-        );
-
         trace_span!("Housekeeping", sync_burn = true).in_scope(|| {
             // TODO: Burn really should implement +=
-            if self.iter > self.config.warmup_steps {
+            if self.iter > self.config.refine_start_iter {
                 // Get the xy gradient norm from the dummy tensor.
                 let xys_grad = Tensor::from_inner(
                     splats
@@ -315,8 +314,9 @@ where
                         .expect("XY gradients need to be calculated."),
                 );
 
-                let aux = &auxes[0];
-                let gs_ids = Tensor::from_primitive(aux.global_from_compact_gid.clone());
+                let aux = auxes[0].clone();
+                let gs_ids = Tensor::from_primitive(aux.global_from_compact_gid);
+                let radii = Tensor::from_primitive(TensorPrimitive::Float(aux.radii));
 
                 let [_, h, w, _] = pred_images.dims();
                 let device = batch.gt_images.device();
@@ -329,17 +329,22 @@ where
                 let num_vis = Tensor::from_primitive(aux.num_visible.clone());
                 let valid = Tensor::arange(0..splats.num_splats() as i64, &device).lower(num_vis);
 
+                self.grad_2d_accum =
+                    self.grad_2d_accum
+                        .clone()
+                        .select_assign(0, gs_ids.clone(), xys_grad_norm);
+
                 self.xy_grad_counts =
                     self.xy_grad_counts
                         .clone()
                         .select_assign(0, gs_ids.clone(), valid.int());
 
-                self.grad_2d_accum = self.grad_2d_accum.clone() + xys_grad_norm;
+                let radii_norm = radii / (w.max(h) as f32);
+                self.max_radii = self.max_radii.clone().max_pair(radii_norm);
             }
         });
 
-        let post_step_splat = trace_span!("Optimizer step", sync_burn = true).in_scope(|| {
-            let mut splats = splats.clone();
+        splats = trace_span!("Optimizer step", sync_burn = true).in_scope(|| {
             let grad_means = GradientsParams::from_params(&mut grads, &splats, &[splats.means.id]);
             splats = self.optim.step(lr_mean, splats, grad_means);
 
@@ -385,18 +390,16 @@ where
 
         let mut refine_stats = None;
 
-        let do_refine = self.iter < self.config.max_refine_step
-            && self.iter >= self.config.warmup_steps
+        let do_refine = self.iter < self.config.refine_stop_iter
+            && self.iter >= self.config.refine_start_iter
             && self.iter % self.config.refine_every == 1;
 
-        splats = if !do_refine {
+        if do_refine {
             // If not refining, update splat to step with gradients applied.
-            post_step_splat
-        } else {
-            let (splats, refine) = self.refine_splats(splats, post_step_splat).await;
+            let (refined_splats, refine) = self.refine_splats(splats, batch.scene_extent).await;
             refine_stats = Some(refine);
-            splats
-        };
+            splats = refined_splats;
+        }
 
         self.iter += 1;
 
@@ -420,21 +423,23 @@ where
     async fn refine_splats(
         &mut self,
         splats: Splats<B>,
-        post_step_splat: Splats<B>,
+        scene_extent: f32,
     ) -> (Splats<B>, RefineStats) {
+        let mut record = self.optim.to_record();
+
+        let mut splats = splats;
+
         let device = splats.means.device();
-        let mut splats_pre_step = splats;
-        let mut splats_post_step = post_step_splat.clone();
 
         // Otherwise, do refinement, but do the split/clone on gaussians with no grads applied.
         let grads = self.grad_2d_accum.clone() / self.xy_grad_counts.clone().clamp_min(1).float();
 
-        let big_grad_mask = grads.greater_equal_elem(self.config.densify_grad_thresh);
-        let split_clone_size_mask = splats_post_step
+        let is_grad_high = grads.greater_equal_elem(self.config.densify_grad_thresh);
+        let split_clone_size_mask = splats
             .scales()
             .max_dim(1)
             .squeeze(1)
-            .lower_elem(self.config.densify_size_thresh);
+            .lower_elem(self.config.densify_size_threshold * scene_extent);
 
         let mut append_means = vec![];
         let mut append_rots = vec![];
@@ -442,107 +447,85 @@ where
         let mut append_opac = vec![];
         let mut append_scales = vec![];
 
-        let clone_inds = Tensor::stack::<2>(
-            vec![split_clone_size_mask.clone(), big_grad_mask.clone()],
-            1,
-        )
-        .all_dim(1)
-        .squeeze::<1>(1)
-        .argwhere_async()
-        .await;
+        let clone_inds =
+            Tensor::stack::<2>(vec![is_grad_high.clone(), split_clone_size_mask.clone()], 1)
+                .all_dim(1)
+                .squeeze::<1>(1)
+                .argwhere_async()
+                .await;
 
         // Clone splats
         let clone_count = clone_inds.dims()[0];
         if clone_count > 0 {
             let clone_inds = clone_inds.squeeze(1);
-            append_means.push(splats_pre_step.means.val().select(0, clone_inds.clone()));
-            append_rots.push(splats_pre_step.rotation.val().select(0, clone_inds.clone()));
-            append_coeffs.push(
-                splats_pre_step
-                    .sh_coeffs
-                    .val()
-                    .select(0, clone_inds.clone()),
-            );
-            append_opac.push(
-                splats_pre_step
-                    .raw_opacity
-                    .val()
-                    .select(0, clone_inds.clone()),
-            );
-            append_scales.push(
-                splats_pre_step
-                    .log_scales
-                    .val()
-                    .select(0, clone_inds.clone()),
-            );
+            let cur_rots = splats.rotation.val().select(0, clone_inds.clone());
+            let cur_scale = splats.log_scales.val().select(0, clone_inds.clone());
+            append_means.push(splats.means.val().select(0, clone_inds.clone()));
+            append_rots.push(cur_rots);
+            append_coeffs.push(splats.sh_coeffs.val().select(0, clone_inds.clone()));
+            append_opac.push(splats.raw_opacity.val().select(0, clone_inds.clone()));
+            append_scales.push(cur_scale);
         }
 
         // Split splats.
         let split_mask =
-            Tensor::stack::<2>(vec![split_clone_size_mask.bool_not(), big_grad_mask], 1).all_dim(1);
-        let split_inds = split_mask.clone().squeeze::<1>(1).argwhere_async().await;
+            Tensor::stack::<2>(vec![is_grad_high, split_clone_size_mask.bool_not()], 1)
+                .all_dim(1)
+                .squeeze(1);
+
+        let radii_grow = self
+            .max_radii
+            .clone()
+            .greater_elem(self.config.densify_radius_threshold);
+        let split_mask = Tensor::stack::<2>(vec![split_mask, radii_grow], 1)
+            .any_dim(1)
+            .squeeze(1);
+
+        let split_inds = split_mask.clone().argwhere_async().await;
         let split_count = split_inds.dims()[0];
+
         if split_count > 0 {
             let split_inds = split_inds.squeeze(1);
 
-            // Some parts can be straightforwardly copied to the new splats.
-            let cur_coeff = splats_post_step
-                .sh_coeffs
-                .val()
-                .select(0, split_inds.clone());
-            let cur_raw_opac = splats_post_step
-                .raw_opacity
-                .val()
-                .select(0, split_inds.clone());
-            let cur_rots = splats_post_step
-                .rotation
-                .val()
-                .select(0, split_inds.clone());
-            append_rots.push(cur_rots.clone());
-            append_coeffs.push(cur_coeff.clone());
-            append_opac.push(cur_raw_opac);
-
-            // Change current scale to be lower.
-            let cur_scale = splats_post_step.scales().select(0, split_inds.clone());
-            Splats::map_param(&mut splats_post_step.log_scales, |m| {
-                let div_scales = Tensor::zeros_like(&m).select_assign(
-                    0,
-                    split_inds.clone(),
-                    (cur_scale.clone() / 1.6).log(),
-                );
-                m.mask_where(split_mask.clone(), div_scales)
-            });
-            // Append newer smaller scales.
-            append_scales.push((cur_scale.clone() / 1.6).log());
-
-            // Sample new position for splits.
-            let cur_means = splats_pre_step.means.val().select(0, split_inds.clone());
-            let samples = quaternion_vec_multiply(
-                cur_rots.clone(),
-                Tensor::random([split_count, 3], Distribution::Normal(0.0, 0.5), &device)
-                    * cur_scale.clone(),
-            );
-            // Assign new means to current values.
-            Splats::map_param(&mut splats_pre_step.means, |m| {
-                let offset_means = Tensor::zeros_like(&m).select_assign(
-                    0,
-                    split_inds.clone(),
-                    cur_means.clone() - samples.clone(),
-                );
-                m.mask_where(split_mask.clone(), offset_means)
-            });
-
-            // Append new means with offset sample.
-            let samples_new = quaternion_vec_multiply(
-                cur_rots.clone(),
-                Tensor::random([split_count, 3], Distribution::Normal(0.0, 0.5), &device)
-                    * cur_scale,
-            );
-            append_means.push(cur_means.clone() + samples_new);
+            for _ in 0..2 {
+                // Some parts can be straightforwardly copied to the new splats.
+                let cur_means = splats.means.val().select(0, split_inds.clone());
+                let cur_coeff = splats.sh_coeffs.val().select(0, split_inds.clone());
+                let cur_raw_opac = splats.raw_opacity.val().select(0, split_inds.clone());
+                let cur_rots = splats.rotation.val().select(0, split_inds.clone());
+                let cur_scale = splats.log_scales.val().select(0, split_inds.clone());
+
+                let samples = quaternion_vec_multiply(
+                    cur_rots.clone(),
+                    Tensor::random([split_count, 3], Distribution::Normal(0.0, 1.0), &device),
+                ) * cur_scale.clone().exp();
+
+                append_means.push(cur_means.clone() + samples);
+                append_rots.push(cur_rots.clone());
+                append_scales.push(cur_scale - 1.6f32.ln());
+                append_coeffs.push(cur_coeff.clone());
+                append_opac.push(cur_raw_opac.clone());
+            }
         }
 
-        // Do processing on splat post step.
-        let mut splats = post_step_splat;
+        prune_points(&mut splats, &mut record, split_mask).await;
+
+        // Do some more processing. Important to do this last as otherwise you might mess up the correspondence
+        // of gradient <-> splat.
+        let start_count = splats.num_splats();
+        // Remove barely visible gaussians.
+        let alpha_mask = splats.opacity().lower_elem(self.config.cull_opacity);
+        prune_points(&mut splats, &mut record, alpha_mask).await;
+        let alpha_pruned = start_count - splats.num_splats();
+
+        // Delete Gaussians with too large of a radius in world-units.
+        let scale_mask = splats
+            .scales()
+            .max_dim(1)
+            .squeeze(1)
+            .greater_elem(self.config.cull_scale3d_percentage_threshold * scene_extent);
+        prune_points(&mut splats, &mut record, scale_mask).await;
+        let scale_pruned = start_count - splats.num_splats();
 
         if !append_means.is_empty() {
             let append_means = Tensor::cat(append_means, 0);
@@ -553,6 +536,7 @@ where
 
             concat_splats(
                 &mut splats,
+                &mut record,
                 append_means,
                 append_rots,
                 append_coeffs,
@@ -561,36 +545,14 @@ where
             );
         }
 
-        // Do some more processing. Important to do this last as otherwise you might mess up the correspondence
-        // of gradient <-> splat.
-        let start_count = splats.num_splats();
-
-        // Remove barely visible gaussians.
-        let alpha_mask = splats.opacity().lower_elem(self.config.cull_alpha_thresh);
-        prune_points(&mut splats, alpha_mask).await;
-
-        let alpha_pruned = start_count - splats.num_splats();
-
-        // Delete Gaussians with too large of a radius in world-units.
-        let scale_mask = splats
-            .scales()
-            .max_dim(1)
-            .squeeze(1)
-            .greater_elem(self.config.cull_scale_thresh);
-        prune_points(&mut splats, scale_mask).await;
-
-        let scale_pruned = start_count - splats.num_splats();
-
         let refine_step = self.iter / self.config.refine_every;
         if refine_step % self.config.reset_alpha_every_refine == 0 {
-            self.reset_opacity(&mut splats);
+            self.reset_opacity(&mut splats, &mut record);
         }
 
         // Stats don't line up anymore so have to reset them.
         self.reset_stats(splats.num_splats(), &device);
-
-        // TODO: Want to do state surgery and keep momenta for splats.
-        self.optim = self.opt_config.init();
+        self.optim = self.optim.clone().load_record(record);
 
         let stats = RefineStats {
             num_split: split_count,
@@ -603,11 +565,30 @@ where
     }
 }
 
+fn map_param<B: AutodiffBackend, const D: usize>(
+    param: &mut Param<Tensor<B, D>>,
+    record: &mut HashMap<ParamId, AdaptorRecord<Adam, B>>,
+    map_param: impl Fn(Tensor<B, D>) -> Tensor<B, D>,
+    map_opt: impl Fn(Tensor<B::InnerBackend, D>) -> Tensor<B::InnerBackend, D>,
+) {
+    Splats::map_param(param, map_param);
+    let mut state: AdamState<_, D> = record[&param.id].clone().into_state();
+    state.momentum.moment_1 = map_opt(state.momentum.moment_1);
+    state.momentum.moment_2 = map_opt(state.momentum.moment_2);
+    record.insert(param.id, AdaptorRecord::from_state(state));
+}
+
 // Prunes points based on the given mask.
 //
 // Args:
 //   mask: bool[n]. If True, prune this Gaussian.
-pub async fn prune_points<B: AutodiffBackend>(splats: &mut Splats<B>, prune: Tensor<B, 1, Bool>) {
+pub async fn prune_points<B: AutodiffBackend>(
+    splats: &mut Splats<B>,
+    record: &mut HashMap<ParamId, AdaptorRecord<Adam, B>>,
+    prune: Tensor<B, 1, Bool>,
+) {
+    assert!(prune.dims()[0] == splats.num_splats());
+
     // bool[n]. If True, delete these Gaussians.
     let prune_count = prune.dims()[0];
 
@@ -620,50 +601,104 @@ pub async fn prune_points<B: AutodiffBackend>(splats: &mut Splats<B>, prune: Ten
     let new_points = valid_inds.dims()[0];
 
     if new_points < start_splats {
-        splats.means = splats
-            .means
-            .clone()
-            .map(|x| Tensor::from_inner(x.select(0, valid_inds.clone()).inner()).require_grad());
-        splats.sh_coeffs = splats
-            .sh_coeffs
-            .clone()
-            .map(|x| Tensor::from_inner(x.select(0, valid_inds.clone()).inner()).require_grad());
-        splats.rotation = splats
-            .rotation
-            .clone()
-            .map(|x| Tensor::from_inner(x.select(0, valid_inds.clone()).inner()).require_grad());
-        splats.raw_opacity = splats
-            .raw_opacity
-            .clone()
-            .map(|x| Tensor::from_inner(x.select(0, valid_inds.clone()).inner()).require_grad());
-        splats.log_scales = splats
-            .log_scales
-            .clone()
-            .map(|x| Tensor::from_inner(x.select(0, valid_inds.clone()).inner()).require_grad());
+        map_param(
+            &mut splats.means,
+            record,
+            |x| x.select(0, valid_inds.clone()),
+            |x| x.select(0, valid_inds.clone().inner()),
+        );
+        map_param(
+            &mut splats.sh_coeffs,
+            record,
+            |x| x.select(0, valid_inds.clone()),
+            |x| x.select(0, valid_inds.clone().inner()),
+        );
+        map_param(
+            &mut splats.rotation,
+            record,
+            |x| x.select(0, valid_inds.clone()),
+            |x| x.select(0, valid_inds.clone().inner()),
+        );
+        map_param(
+            &mut splats.raw_opacity,
+            record,
+            |x| x.select(0, valid_inds.clone()),
+            |x| x.select(0, valid_inds.clone().inner()),
+        );
+        map_param(
+            &mut splats.log_scales,
+            record,
+            |x| x.select(0, valid_inds.clone()),
+            |x| x.select(0, valid_inds.clone().inner()),
+        );
     }
 }
 
 pub fn concat_splats<B: AutodiffBackend>(
     splats: &mut Splats<B>,
+    record: &mut HashMap<ParamId, AdaptorRecord<Adam, B>>,
     means: Tensor<B, 2>,
     rotations: Tensor<B, 2>,
     sh_coeffs: Tensor<B, 3>,
-    raw_opacities: Tensor<B, 1>,
+    raw_opac: Tensor<B, 1>,
     log_scales: Tensor<B, 2>,
 ) {
-    Splats::map_param(&mut splats.means, |x| {
-        Tensor::cat(vec![x, means.clone()], 0)
-    });
-    Splats::map_param(&mut splats.rotation, |x| {
-        Tensor::cat(vec![x, rotations.clone()], 0)
-    });
-    Splats::map_param(&mut splats.sh_coeffs, |x| {
-        Tensor::cat(vec![x, sh_coeffs.clone()], 0)
-    });
-    Splats::map_param(&mut splats.raw_opacity, |x| {
-        Tensor::cat(vec![x, raw_opacities.clone()], 0)
-    });
-    Splats::map_param(&mut splats.log_scales, |x| {
-        Tensor::cat(vec![x, log_scales.clone()], 0)
-    });
+    // Concat
+    map_param(
+        &mut splats.means,
+        record,
+        |x| Tensor::cat(vec![x, means.clone()], 0),
+        |x| Tensor::cat(vec![x, Tensor::zeros_like(&means.clone().inner())], 0),
+    );
+    map_param(
+        &mut splats.rotation,
+        record,
+        |x| Tensor::cat(vec![x, rotations.clone()], 0),
+        |x| Tensor::cat(vec![x, Tensor::zeros_like(&rotations.clone().inner())], 0),
+    );
+    map_param(
+        &mut splats.sh_coeffs,
+        record,
+        |x| Tensor::cat(vec![x, sh_coeffs.clone()], 0),
+        |x| Tensor::cat(vec![x, Tensor::zeros_like(&sh_coeffs.clone().inner())], 0),
+    );
+    map_param(
+        &mut splats.raw_opacity,
+        record,
+        |x| Tensor::cat(vec![x, raw_opac.clone()], 0),
+        |x| Tensor::cat(vec![x, Tensor::zeros_like(&raw_opac.clone().inner())], 0),
+    );
+    map_param(
+        &mut splats.log_scales,
+        record,
+        |x| Tensor::cat(vec![x, log_scales.clone()], 0),
+        |x| Tensor::cat(vec![x, Tensor::zeros_like(&log_scales.clone().inner())], 0),
+    );
+}
+
+#[cfg(test)]
+mod tests {
+    use burn::{
+        backend::{wgpu::WgpuDevice, Wgpu},
+        tensor::Tensor,
+    };
+    use glam::Quat;
+
+    use super::quaternion_vec_multiply;
+
+    #[test]
+    fn test_quat_multiply() {
+        let quat = Quat::from_euler(glam::EulerRot::XYZ, 0.2, 0.2, 0.3);
+        let vec = glam::vec3(0.5, 0.7, 0.1);
+        let result_ref = quat * vec;
+
+        let device = WgpuDevice::DefaultDevice;
+        let quaternions = Tensor::<Wgpu, 1>::from_floats([quat.w, quat.x, quat.y, quat.z], &device)
+            .reshape([1, 4]);
+        let vecs = Tensor::<Wgpu, 1>::from_floats([vec.x, vec.y, vec.z], &device).reshape([1, 3]);
+        let result = quaternion_vec_multiply(quaternions, vecs);
+        let result: Vec<f32> = result.into_data().to_vec().unwrap();
+        let result = glam::vec3(result[0], result[1], result[2]);
+        assert!((result_ref - result).length() < 1e-7);
+    }
 }
diff --git a/crates/brush-viewer/src/panels/load_data.rs b/crates/brush-viewer/src/panels/load_data.rs
index 8d5cd71e..c5b565ab 100644
--- a/crates/brush-viewer/src/panels/load_data.rs
+++ b/crates/brush-viewer/src/panels/load_data.rs
@@ -10,8 +10,7 @@ enum Quality {
 
 pub(crate) struct LoadDataPanel {
     load_args: LoadDatasetArgs,
-
-    sh_degree: u32,
+    init_args: LoadInitArgs,
     quality: Quality,
     url: String,
 }
@@ -22,13 +21,10 @@ impl LoadDataPanel {
             // Super high resolutions are a bit sketchy. Limit to at least
             // some size.
             load_args: LoadDatasetArgs {
-                max_frames: None,
                 max_resolution: Some(1920),
-                eval_split_every: None,
-                subsample_frames: None,
-                subsample_points: None,
+                ..Default::default()
             },
-            sh_degree: 3,
+            init_args: LoadInitArgs::default(),
             quality: Quality::Normal,
             url: "splat.com/example.ply".to_owned(),
         }
@@ -54,17 +50,13 @@ impl ViewerPanel for LoadDataPanel {
             ui.add_space(10.0);
 
             if file || url {
-                let load_init_args = LoadInitArgs {
-                    sh_degree: self.sh_degree,
-                };
-
                 let mut config = TrainConfig::default();
                 if matches!(self.quality, Quality::Low) {
                     config = config
                         .with_densify_grad_thresh(0.0003)
                         .with_refine_every(150)
                         .with_ssim_weight(0.0)
-                        .with_cull_alpha_thresh(0.01);
+                        .with_cull_opacity(0.01);
                 }
 
                 let source = if file {
@@ -72,14 +64,19 @@ impl ViewerPanel for LoadDataPanel {
                 } else {
                     crate::viewer::DataSource::Url(self.url.to_string())
                 };
-                context.start_data_load(source, self.load_args.clone(), load_init_args, config);
+                context.start_data_load(
+                    source,
+                    self.load_args.clone(),
+                    self.init_args.clone(),
+                    config,
+                );
             }
 
             ui.add_space(10.0);
             ui.heading("Train settings");
 
             ui.label("Spherical Harmonics Degree:");
-            ui.add(Slider::new(&mut self.sh_degree, 0..=4));
+            ui.add(Slider::new(&mut self.init_args.sh_degree, 0..=4));
 
             ui.horizontal(|ui| {
                 ui.label("Quality:");
diff --git a/crates/brush-viewer/src/train_loop.rs b/crates/brush-viewer/src/train_loop.rs
index 9c70b54a..03c10921 100644
--- a/crates/brush-viewer/src/train_loop.rs
+++ b/crates/brush-viewer/src/train_loop.rs
@@ -59,7 +59,7 @@ pub(crate) fn train_loop<T: AsyncRead + Unpin + 'static>(
         // Read initial splats if any.
         while let Some(message) = splat_stream.next().await {
             let message = message?;
-            let splats = message.splats.with_min_sh_degree(load_init_args.sh_degree);
+            let splats = message.splats.with_sh_degree(load_init_args.sh_degree);
             let msg = ProcessMessage::ViewSplats {
                 up_axis: message.meta.up_axis,
                 splats: Box::new(splats.valid()),
@@ -88,11 +88,13 @@ pub(crate) fn train_loop<T: AsyncRead + Unpin + 'static>(
             splats
         } else {
             // By default, spawn the splats in bounds.
-            let bounds = dataset.train.bounds(0.0, 0.0);
+            let bounds = dataset.train.bounds();
             let bounds_extent = bounds.extent.length();
             // Arbitrarly assume area of interest is 0.2 - 0.75 of scene bounds.
             // Somewhat specific to the blender scenes
-            let adjusted_bounds = dataset.train.bounds(bounds_extent * 0.25, bounds_extent);
+            let adjusted_bounds = dataset
+                .train
+                .adjusted_bounds(bounds_extent * 0.25, bounds_extent);
 
             let config = RandomSplatsConfig::new().with_sh_degree(load_init_args.sh_degree);
             Splats::from_random_config(config, adjusted_bounds, &mut rng, &device)
diff --git a/crates/brush-viewer/src/viewer.rs b/crates/brush-viewer/src/viewer.rs
index ff9d3297..9d838c46 100644
--- a/crates/brush-viewer/src/viewer.rs
+++ b/crates/brush-viewer/src/viewer.rs
@@ -320,10 +320,7 @@ impl ViewerContext {
 
         // Copy the camera, mostly to copy over the intrinsics and such.
         self.controls.focus = transform.translation
-            + transform.matrix3
-                * Vec3A::Z
-                * self.dataset.train.bounds(0.0, 0.0).extent.length()
-                * 0.5;
+            + transform.matrix3 * Vec3A::Z * self.dataset.train.bounds().extent.length() * 0.5;
         self.controls.dirty = true;
         self.camera = cam.clone();
     }
diff --git a/crates/train-2d/src/main.rs b/crates/train-2d/src/main.rs
index f8346e3d..316e2e45 100644
--- a/crates/train-2d/src/main.rs
+++ b/crates/train-2d/src/main.rs
@@ -135,8 +135,8 @@ impl App {
                 .load_texture("nearest_view_tex", color_img, TextureOptions::default());
 
         let config = TrainConfig::new(ExponentialLrSchedulerConfig::new(lr_max, decay))
-            .with_max_refine_step(u32::MAX) // Just keep refining
-            .with_warmup_steps(100) // Don't really need a warmup for simple 2D
+            .with_refine_start_iter(100) // Don't really need a warmup for simple 2D
+            .with_refine_stop_iter(u32::MAX) // Just keep refining
             .with_reset_alpha_every_refine(u32::MAX); // Don't use alpha reset.
 
         spawn_train_loop(view.clone(), config, device, cc.egui_ctx.clone(), sender);
diff --git a/deny.toml b/deny.toml
index ed2018ea..5c1569b8 100644
--- a/deny.toml
+++ b/deny.toml
@@ -38,24 +38,16 @@ skip = [
     { name = 'cfg_aliases' },
     { name = 'syn' },
     { name = 'itertools' },
+    { name = 'derive-new' },
+
+    #
+    { name = 'thiserror' },
+    { name = 'thiserror-impl' },
+    { name = 'unicode-width' },
 
     # Rfd problematic deps.
-    { name = 'futures-lite' },      # Duplicate in rfd.
-    { name = 'fastrand' },          # Duplicate in rfd.
-    { name = 'event-listener' },    # Duplicate in rfd.
-    { name = 'async-io' },          # Duplicate in rfd.
-    { name = 'async-lock' },        # Duplicate in rfd.
-    { name = 'async-fs' },          # Duplicate in rfd.
-    { name = 'polling' },           # Duplicate in rfd.
-    { name = 'memoffset' },         # Duplicate in rfd.
-    { name = 'linux-raw-sys' },     # Duplicate in rfd.
-    { name = 'toml_edit' },         # Duplicate in rfd.
-    { name = 'rustix' },            # Duplicate in rfd.
     { name = 'raw-window-handle' }, # Old version from rfd.
 
-    { name = 'async-process' }, # RFD & async-std
-    { name = 'async-channel' }, # async-std
-
     # Custom ply fork.
     { name = "ply-rs" },         # Custom fork & rerun version.
     { name = 'peg' },            # Custom ply fork & rerun version.
@@ -63,12 +55,8 @@ skip = [
     { name = 'peg-runtime' },    # Custom ply fork & rerun version.
     { name = 'cargo_metadata' }, # Duplicate from custom ply + rerun
 
-    # Rerun
-    { name = 'fixed' }, # Old version iin rerun
-
     # Other
     { name = "core-foundation" }, # Egui has two?
-    { name = 'bincode' },         # Burn uses rc2, puffin 1.3.1
     # { name = "web-time" },            # Old version from egui & rerun.
     { name = 'windows_x86_64_msvc' }, # Old version from winit.
     { name = 'windows_x86_64_gnu' },  # Old version from winit.
@@ -98,9 +86,10 @@ allow = [
     "MIT",                            # https://tldrlegal.com/license/mit-license
     "MPL-2.0",                        # https://www.mozilla.org/en-US/MPL/2.0/FAQ/ - see Q11. Used by webpki-roots on Linux.
     "OFL-1.1",                        # https://spdx.org/licenses/OFL-1.1.html
-    "Unicode-DFS-2016",               # https://spdx.org/licenses/Unicode-DFS-2016.html
     "Zlib",                           # https://tldrlegal.com/license/zlib-libpng-license-(zlib)
     "Unlicense",                      # https://www.tldrlegal.com/license/unlicense
+    "Unicode-3.0",
+    "OpenSSL",
 ]
 exceptions = []