Upgrade to burn 0.21.0-pre.3 and consolidate workspace#96
Merged
antimora merged 13 commits intotracel-ai:mainfrom Apr 16, 2026
Merged
Upgrade to burn 0.21.0-pre.3 and consolidate workspace#96antimora merged 13 commits intotracel-ai:mainfrom
antimora merged 13 commits intotracel-ai:mainfrom
Conversation
Bump burn and burn-store across all model crates and adapt to the 0.20 -> 0.21 API changes: - PaddingConfig2d::Explicit widened from (h, w) to (top, left, bottom, right) in resnet-burn, mobilenetv2-burn, yolox-burn. - Shape.dims is no longer a public field; use .dims::<D>() in bert-burn. - TransformerEncoderConfig gained required activation and layer_norm_eps fields; PositionWiseFeedForward's gelu field is now an ActivationRecord. - burn-import ModelGen::embed_states(bool) replaced with load_strategy(LoadStrategy) in squeezenet-burn. - albert-burn drops its git rev pin in favor of the published release.
burn-import::onnx::ModelGen is deprecated in 0.21 in favor of burn_onnx::ModelGen. Depend on burn-onnx directly in squeezenet-burn, which is the only crate that used the ONNX code generator.
- tokenizers 0.15/0.19 -> 0.22 - hf-hub 0.3/0.4 -> 0.5 - dirs 5 -> 6 (unified across crates) - image 0.24 -> 0.25 - itertools 0.12 -> 0.14 - rand 0.8 -> 0.9 (fixes distributions -> distr::weighted rename in sampling) - tiktoken-rs 0.5 -> 0.11 (Rank is now u32; decode takes &[u32]) - candle-core 0.8 -> 0.9 in bert-burn; dropped from squeezenet-burn (unused after the burn-onnx switch) - serde/tokio/libm/clap/csv/tar/etc. relaxed to major-only constraints so cargo picks the newest compatible release - bert-burn: drop the cubecl-runtime workaround; burn 0.21 no longer needs it llama-burn keeps rustc-hash 1.1 because tiktoken-rs still depends on that major version.
burn-flex is a pure-Rust CPU backend with SIMD acceleration and fewer build dependencies than burn-ndarray. Swap the CPU backend used by examples, tests, and benches across all model crates. - Rename the "ndarray" feature to "flex" and add burn-flex as an optional dependency (library crates) or dev-dependency (example-only crates). - Replace `NdArray<f32>` with `burn_flex::Flex` (Flex fixes FloatElem to f32, so the generic parameter is dropped). - Replace `NdArrayDevice::Cpu` with `burn_flex::FlexDevice`.
Add a top-level workspace with shared package metadata and shared dep
versions. Single source of truth for burn / serde / tokio / image /
tokenizers / etc., one Cargo.lock, one target dir.
- Resolver 3 (Rust 2024 edition) so members on edition 2021 and 2024
can coexist without feature unification surprises.
- All non-burn deps move to [workspace.dependencies] with major-only
pins; per-crate manifests inherit via { workspace = true } and add
their own feature flags.
- Shared license = "MIT OR Apache-2.0".
Flatten resnet-burn: the inner resnet-burn/resnet/ crate moves up to
resnet-burn/. The old inner [workspace] is gone, so the two examples
join the top-level workspace alongside the model crate. Example path
deps update from "../../resnet" to "../..".
Drop all backend feature flags (tch-cpu, tch-gpu, wgpu, cuda, fusion,
ndarray, f16) and keep only Flex. Burn is removing the B-generic, so
each example/test/bench targets Flex directly via burn-flex (now a
dev-dependency, not optional). The fan-out cfg modules in bert-burn
and llama-burn examples collapse to a single launch::<Flex>(FlexDevice).
llama-burn loses two dead test modules (transformer.rs and llama.rs)
that were gated behind the now-removed cuda/tch-gpu features and had
bit-rotted against the current burn API.
- bert-burn loader: ConstantRecord -> EmptyRecord (the old type alias is deprecated in burn 0.21 because the name implied data persistence it never provided). - minilm-burn bench: criterion::black_box -> std::hint::black_box (the criterion re-export is deprecated in 0.6). Workspace `cargo check --all-targets` is now warning-free.
Now that the repo is one workspace producing multiple binaries (model examples, training, classification), pin the dep graph by tracking Cargo.lock. CI runs with --locked so a CI failure means the lockfile is out of date, not that an upstream release silently changed under us. Replace the squeezenet-only CI steps with workspace-wide cargo fmt, clippy, build, and test from the repo root. Add Swatinem/rust-cache since the workspace target dir is shared and large.
burn 0.21.0-pre.3 requires rustc 1.92 (per its rust-version field), so the old 1.89.0 matrix entry fails at resolve time.
There was a problem hiding this comment.
Pull request overview
Upgrades the repository to the burn 0.21.0-pre.3 ecosystem and consolidates the previously separate model crates/examples into a single workspace, standardizing on the burn-flex backend and updating CI to build/test the full workspace with a locked lockfile.
Changes:
- Bump
burn/burn-store/burn-onnx/burn-flexand modernize shared third-party dependencies via workspace-managed versions. - Replace ndarray-based examples/tests with
burn_flex::Flexand update API breakages (e.g.,PaddingConfig2d::Explicit(...)arity, new activation/record APIs). - Consolidate ResNet into a single crate + add pretrained weights metadata and training hooks; update CI to run workspace-wide with
--locked.
Reviewed changes
Copilot reviewed 48 out of 53 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| yolox-burn/src/model/head.rs | Update padding API usage for Burn 0.21 pre. |
| yolox-burn/src/model/bottleneck.rs | Update padding API usage for pooling. |
| yolox-burn/src/model/blocks.rs | Update padding API usage for conv blocks. |
| yolox-burn/examples/inference.rs | Switch example backend from NdArray to Flex. |
| yolox-burn/Cargo.toml | Move deps to workspace + add burn-flex dev-dep. |
| squeezenet-burn/examples/classify.rs | Switch example backend/device from NdArray to Flex. |
| squeezenet-burn/build.rs | Replace burn-import with burn-onnx + load strategy API. |
| squeezenet-burn/Cargo.toml | Move deps to workspace + switch build-dep to burn-onnx. |
| resnet-burn/src/weights.rs | Add pretrained weight metadata + downloader helper. |
| resnet-burn/src/training.rs | Add TrainStep/InferenceStep integration for training. |
| resnet-burn/src/resnet.rs | Update padding API usage for stem layers. |
| resnet-burn/src/lib.rs | New crate root wiring (std/train feature gating). |
| resnet-burn/src/block.rs | Update padding API usage for ResNet blocks. |
| resnet-burn/resnet/Cargo.toml | Remove nested crate as part of consolidation. |
| resnet-burn/examples/inference/examples/inference.rs | Switch inference example backend from NdArray to Flex. |
| resnet-burn/examples/inference/Cargo.toml | Update example to depend on consolidated resnet-burn crate + Flex. |
| resnet-burn/examples/finetune/examples/finetune.rs | Remove multi-backend feature wiring; standardize on Flex. |
| resnet-burn/examples/finetune/Cargo.toml | Update deps to workspace + add burn-flex. |
| resnet-burn/Cargo.toml | Convert from workspace root to a single published crate in the top-level workspace. |
| mobilenetv2-burn/src/model/inverted_residual.rs | Update padding API usage. |
| mobilenetv2-burn/src/model/conv_norm.rs | Update padding API usage (explicit 4-side padding). |
| mobilenetv2-burn/examples/inference.rs | Switch inference example backend from NdArray to Flex. |
| mobilenetv2-burn/Cargo.toml | Move deps to workspace + add burn-flex dev-dep. |
| minilm-burn/tests/integration_test.rs | Switch ignored integration tests from NdArray feature gating to Flex. |
| minilm-burn/src/pooling.rs | Run unit tests with Flex backend (no ndarray feature gating). |
| minilm-burn/src/lib.rs | Update docs example backend from NdArray to Flex. |
| minilm-burn/examples/inference.rs | Switch example backend from NdArray to Flex. |
| minilm-burn/benches/inference.rs | Simplify benchmarks to a single Flex backend; update black_box usage. |
| minilm-burn/Cargo.toml | Remove backend-selection features; use workspace deps; add burn-flex dev-dep. |
| llama-burn/src/transformer.rs | Adjust test gating + update tolerance assertion API. |
| llama-burn/src/tokenizer/tiktoken.rs | Update token id/rank types + decode API for newer tiktoken-rs. |
| llama-burn/src/sampling.rs | Update rand 0.9 import paths for distributions. |
| llama-burn/src/llama.rs | Remove backend-gated tests block (consolidation). |
| llama-burn/src/lib.rs | Standardize TestBackend to Flex for unit tests. |
| llama-burn/examples/test_tiny.rs | Remove backend feature selection; run with Flex. |
| llama-burn/examples/chat.rs | Remove backend feature selection; run with Flex. |
| llama-burn/Cargo.toml | Remove backend-selection features; move deps to workspace; add burn-flex dev-dep. |
| bert-burn/src/model.rs | Update transformer config wiring for new activation/layer-norm fields. |
| bert-burn/src/loader.rs | Update record types (ConstantRecord→EmptyRecord) + explicit activation record. |
| bert-burn/src/embedding.rs | Update tensor shape API usage (dims::<2>()). |
| bert-burn/examples/masked.rs | Remove backend feature wiring; run with Flex. |
| bert-burn/examples/infer-embedding.rs | Remove backend feature wiring; run with Flex. |
| bert-burn/Cargo.toml | Move deps to workspace; drop backend-selection features; add burn-flex dev-dep. |
| albert-burn/tests/integration_test.rs | Switch ignored integration tests from ndarray gating to Flex. |
| albert-burn/src/lib.rs | Update docs example backend from NdArray to Flex. |
| albert-burn/src/encoder.rs | Add clippy allow for many-args constructor. |
| albert-burn/examples/inference.rs | Switch example backend from NdArray to Flex. |
| albert-burn/benches/inference.rs | Simplify benches to a single Flex backend. |
| albert-burn/Cargo.toml | Move deps to workspace; remove backend-selection features; add burn-flex dev-dep. |
| Cargo.toml | Add top-level workspace, members list, and centralized dependency versions. |
| .gitignore | Stop ignoring Cargo.lock (track lockfile) and adjust comments. |
| .github/workflows/rust.yml | Run fmt/clippy/build/test at workspace scope with --locked + caching. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- Silence excessive_precision / needless_range_loop in albert and minilm integration tests; their constants are verbatim Python reference outputs that we want to diff against future regenerations. - squeezenet classify example: allow needless_range_loop on the pixel loop where y,x index three parallel channel arrays. - squeezenet build.rs: create target/<profile>/examples/ before copying the .bpk, since cargo only creates it once an example is actually built and the build script runs first.
The #[ignore] line comments still referenced '--features flex', but minilm-burn no longer defines a flex feature (burn-flex is now a plain dev-dependency after 59be387). Replace with the feature-free 'cargo test -p minilm-burn -- --ignored' invocation. Addresses Copilot review on PR tracel-ai#96.
laggui
reviewed
Apr 16, 2026
Member
laggui
left a comment
laggui
left a comment
There was a problem hiding this comment.
Drop tch / wgpu / cuda / fusion / ndarray / f16 backend flags (anticipating removal of B: Backend generic)
We'll still have to enable the burn flags based on the enabled backends though.
Some of the current examples might run a bit slow on the flex CPU backend 😅
Otherwise LGTM
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
burn0.20.0 → 0.21.0-pre.3burn-importwithburn-onnxburn-ndarraywithburn-flexB: Backendgeneric)Cargo.lock; CI now runs workspace-wide with--lockedTest plan
cargo check --workspace --all-targets— warning-freecargo check --workspace --all-targets --locked— lockfile matches