collective
Browse / github.com/storytold/craft-ai-integrations / patch
Enforce Rust coding standards via Clippy
Enforce Rust coding standards via Clippy Replaces #110 (closed when its head branch was deleted). Strict workspace Clippy lint settings plus the refactors needed to satisfy them. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **Bug Fixes** * Improved tensor shape handling in AI model processing, including cases where singleton dimensions need to expand to match input dimensions. * **Improvements** * Updated generated model components to use more precise imports and consistent handling of tensor expansion. * **Developer Experience** *
- Status
- open
- Score
- 0.00
- Adopters
- 0 other operators
- Author
- dylanmtaylor (imported from pull request #111, open)
- Base
6f7078666a341900987b55585dd9060d6ccd05d1- Size
- +200 −233 in 36 files
- Review
- pending
- Shared
- 2026-10-28
- Id
d0310395f23b16f6914dc810
How it went
No other operator has reported using it yet.
Diff
raw · AGENTS.md Cargo.toml clippy.toml crates/ai-core/tests/draft_provenance.rs crates/ai-core/tests/mock_conformance.rs crates/ai-engine-burn/src/model/decoder/mod.rs crates/ai-engine-burn/src/model/decoder/part_01.rs crates/ai-engine-burn/src/model/decoder/part_01/forward.rs crates/ai-engine-burn/src/model/decoder/part_02.rs crates/ai-engine-burn/src/model/decoder/part_03.rs crates/ai-engine-burn/src/model/decoder/part_03/forward.rs crates/ai-engine-burn/src/model/decoder/part_04.rs crates/ai-engine-burn/src/model/decoder/part_04/forward.rs crates/ai-engine-burn/src/model/decoder/part_05.rs crates/ai-engine-burn/src/model/encoder/mod.rs crates/ai-engine-burn/src/model/encoder/part_01.rs crates/ai-engine-burn/src/model/encoder/part_01/forward.rs crates/ai-engine-burn/src/model/encoder/part_02.rs crates/ai-engine-burn/src/model/encoder/part_03.rs crates/ai-engine-burn/src/model/encoder/part_04.rs crates/ai-engine-burn/src/model/encoder/part_05.rs crates/ai-engine-burn/src/model/encoder/part_05/forward.rs crates/ai-engine-burn/src/model/encoder/part_06.rs crates/ai-engine-burn/src/model/encoder/part_06/forward.rs crates/ai-engine-burn/src/model/encoder/part_07.rs crates/ai-engine-burn/src/model/encoder/part_07/forward.rs crates/ai-engine-burn/src/model/encoder/part_08.rs crates/ai-engine-burn/src/model/encoder/part_08/forward.rs crates/ai-engine-burn/src/model/encoder/part_09.rs crates/ai-engine-burn/src/model/encoder/part_09/forward.rs crates/ai-engine-burn/src/model/encoder/part_10.rs crates/ai-engine-burn/src/model/encoder/part_11.rs crates/ai-engine-burn/src/model/encoder/part_12.rs crates/ai-engine-burn/src/model/encoder/part_13.rs crates/ai-engine-burn/src/model/mod.rs tools/burn/generate_sam21.py
@@ -6,6 +6,9 @@ editing. For a developer suggestion about a model, connection, engine or workflo recommendation, a working integration or a user account. Keep credentials and user bindings out of proposals and version control. +Every new workspace crate's `Cargo.toml` must explicitly include `[lints] workspace = true`, and+production crate roots must retain their required crate-level lint attribute.+ The repository's shared crates own the contracts, adapters, catalogue, consent and client. Each Craft app owns its own interface, document conversion and result storage. Check the live catalogue, route implementation and app adoption before claiming a provider or model is missing@@ -21,6 +21,9 @@ chrono = { version = "0.4.45", default-features = false, features = ["std"] } unsafe_code = "forbid" [workspace.lints.clippy]+cognitive_complexity = "deny"+too_many_arguments = "deny"+wildcard_imports = "deny" dbg_macro = "warn" todo = "warn" @@ -2,3 +2,4 @@ allow-unwrap-in-tests = true allow-expect-in-tests = true allow-panic-in-tests = true allow-indexing-slicing-in-tests = true+cognitive-complexity-threshold = 30@@ -1,5 +1,5 @@ //! Connection-test results belong only to the address they were requested from.-use craft_ai_core::*;+use craft_ai_core::{ConnectionDraft, ModelId, ModelInfo, PRESETS}; /// A draft pointing at `endpoint` that uses no key. fn draft(endpoint: &str) -> ConnectionDraft {@@ -7,7 +7,11 @@ use craft_ai_core::consent::Disclosure; use craft_ai_core::job::{Event, JobState, Output}; use craft_ai_core::mock::{MockProvider, Outcome, Script}; use craft_ai_core::role::{Binding, Role};-use craft_ai_core::*;+use craft_ai_core::{+ AiError, AppId, Auth, Billing, CapabilityId, CapabilitySpec, Connection, ConnectionId,+ ConsentLedger, ErrorKind, Features, InputRole, Limits, Locality, ModelId, ModelInfo, Payload,+ Request, RoleId, SecretRef, StoreFile,+}; fn cap(s: &str) -> CapabilityId { CapabilityId::parse(s).unwrap()@@ -1,13 +1,8 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)]-use burn::nn::LayerNorm;-use burn::nn::LayerNormConfig;-use burn::nn::Linear;-use burn::nn::LinearConfig;-use burn::nn::conv::ConvTranspose2d;-use burn::nn::conv::ConvTranspose2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ mod part_01; use part_01::Submodule1; mod part_02;@@ -1,13 +1,10 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)]-use burn::nn::LayerNorm;-use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig;-use burn::nn::conv::ConvTranspose2d;-use burn::nn::conv::ConvTranspose2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule1 { constant134: burn::module::Param<Tensor<1>>,@@ -1,5 +1,7 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-use super::*;+use super::super::super::broadcast_shape;+use super::{Bool, Int, Module, Submodule1, Tensor, s};+ impl Submodule1 { pub fn forward( &self,@@ -102,17 +104,9 @@ impl Submodule1 { output }; let expand1_out1 = {- let onnx_shape: [i64; 2usize] = shape4_out1;+ let onnx_shape = shape4_out1; let input_dims = div1_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..2usize {- let dim_offset = i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- div1_out1.expand(shape)+ div1_out1.expand(broadcast_shape(onnx_shape, input_dims, 0)) }; let shape5_out1: [i64; 3] = { let axes = &concat4_out1.clone().dims()[0..3];@@ -181,20 +175,12 @@ impl Submodule1 { ) .mask_where(equal1_out1, constantofshape4_out1.clone()); let expand2_out1 = {- let onnx_shape: [i64; 3usize] = TryInto::<[i64; 3usize]>::try_into(+ let onnx_shape = TryInto::<[i64; 3usize]>::try_into( where1_out1.to_data().convert::<i64>().as_slice().unwrap(), ) .unwrap(); let input_dims = reshape1_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..3usize {- let dim_offset = i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- reshape1_out1.expand(shape)+ reshape1_out1.expand(broadcast_shape(onnx_shape, input_dims, 0)) }; let unsqueeze4_out1: Tensor<4, Int> = expand2_out1.unsqueeze_dims::<4>(&[-1]); let constant156_out1 = self.constant156.val();@@ -210,20 +196,12 @@ impl Submodule1 { ) .mask_where(equal2_out1, constantofshape4_out1.clone()); let expand3_out1 = {- let onnx_shape: [i64; 3usize] = TryInto::<[i64; 3usize]>::try_into(+ let onnx_shape = TryInto::<[i64; 3usize]>::try_into( where2_out1.to_data().convert::<i64>().as_slice().unwrap(), ) .unwrap(); let input_dims = reshape2_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..2usize {- let dim_offset = 1usize + i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- reshape2_out1.expand(shape)+ reshape2_out1.expand(broadcast_shape(onnx_shape, input_dims, 1)) }; let unsqueeze5_out1: Tensor<4, Int> = expand3_out1.unsqueeze_dims::<4>(&[-1]); let constant158_out1 = self.constant158.val();@@ -240,20 +218,12 @@ impl Submodule1 { .mask_where(equal3_out1, constantofshape4_out1); let constant150_out1 = self.constant150.val(); let expand4_out1 = {- let onnx_shape: [i64; 3usize] = TryInto::<[i64; 3usize]>::try_into(+ let onnx_shape = TryInto::<[i64; 3usize]>::try_into( where3_out1.to_data().convert::<i64>().as_slice().unwrap(), ) .unwrap(); let input_dims = constant150_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..1usize {- let dim_offset = 2usize + i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- constant150_out1.expand(shape)+ constant150_out1.expand(broadcast_shape(onnx_shape, input_dims, 2)) }; let unsqueeze6_out1: Tensor<4, Int> = expand4_out1.unsqueeze_dims::<4>(&[-1]); let concat6_out1 = burn::tensor::Tensor::cat(@@ -310,17 +280,9 @@ impl Submodule1 { output }; let expand5_out1 = {- let onnx_shape: [i64; 2usize] = shape12_out1;+ let onnx_shape = shape12_out1; let input_dims = div2_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..2usize {- let dim_offset = i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- div2_out1.expand(shape)+ div2_out1.expand(broadcast_shape(onnx_shape, input_dims, 0)) }; let shape13_out1: [i64; 3] = { let axes = &scatternd1_out1.clone().dims()[0..3];@@ -391,20 +353,12 @@ impl Submodule1 { ) .mask_where(equal4_out1, constantofshape7_out1.clone()); let expand6_out1 = {- let onnx_shape: [i64; 3usize] = TryInto::<[i64; 3usize]>::try_into(+ let onnx_shape = TryInto::<[i64; 3usize]>::try_into( where4_out1.to_data().convert::<i64>().as_slice().unwrap(), ) .unwrap(); let input_dims = reshape4_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..3usize {- let dim_offset = i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- reshape4_out1.expand(shape)+ reshape4_out1.expand(broadcast_shape(onnx_shape, input_dims, 0)) }; let unsqueeze7_out1: Tensor<4, Int> = expand6_out1.unsqueeze_dims::<4>(&[-1]); let constant176_out1 = self.constant176.val();@@ -420,20 +374,12 @@ impl Submodule1 { ) .mask_where(equal5_out1, constantofshape7_out1.clone()); let expand7_out1 = {- let onnx_shape: [i64; 3usize] = TryInto::<[i64; 3usize]>::try_into(+ let onnx_shape = TryInto::<[i64; 3usize]>::try_into( where5_out1.to_data().convert::<i64>().as_slice().unwrap(), ) .unwrap(); let input_dims = reshape5_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..2usize {- let dim_offset = 1usize + i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- reshape5_out1.expand(shape)+ reshape5_out1.expand(broadcast_shape(onnx_shape, input_dims, 1)) }; let unsqueeze8_out1: Tensor<4, Int> = expand7_out1.unsqueeze_dims::<4>(&[-1]); let constant178_out1 = self.constant178.val();@@ -450,20 +396,12 @@ impl Submodule1 { .mask_where(equal6_out1, constantofshape7_out1); let constant170_out1 = self.constant170.val(); let expand8_out1 = {- let onnx_shape: [i64; 3usize] = TryInto::<[i64; 3usize]>::try_into(+ let onnx_shape = TryInto::<[i64; 3usize]>::try_into( where6_out1.to_data().convert::<i64>().as_slice().unwrap(), ) .unwrap(); let input_dims = constant170_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..1usize {- let dim_offset = 2usize + i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- constant170_out1.expand(shape)+ constant170_out1.expand(broadcast_shape(onnx_shape, input_dims, 2)) }; let unsqueeze9_out1: Tensor<4, Int> = expand8_out1.unsqueeze_dims::<4>(&[-1]); let concat8_out1 = burn::tensor::Tensor::cat(@@ -625,15 +563,7 @@ impl Submodule1 { ) .unwrap(); let input_dims = constant91_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..4usize {- let dim_offset = i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- constant91_out1.expand(shape)+ constant91_out1.expand(broadcast_shape(onnx_shape, input_dims, 0)) }; let shape22_out1: [i64; 3] = { let axes = &concat11_out1.clone().dims()[0..3];@@ -677,20 +607,12 @@ impl Submodule1 { .mask_where(equal13_out1, constantofshape12_out1); let constant92_out1 = self.constant92.val(); let expand10_out1 = {- let onnx_shape: [i64; 3usize] = TryInto::<[i64; 3usize]>::try_into(+ let onnx_shape = TryInto::<[i64; 3usize]>::try_into( where13_out1.to_data().convert::<i64>().as_slice().unwrap(), ) .unwrap(); let input_dims = constant92_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..3usize {- let dim_offset = i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- constant92_out1.expand(shape)+ constant92_out1.expand(broadcast_shape(onnx_shape, input_dims, 0)) }; let concat14_out1 = burn::tensor::Tensor::cat([expand10_out1, concat11_out1].into(), 1); let add11_out1 = image_embed.add(expand9_out1);@@ -1,13 +1,10 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)]-use burn::nn::LayerNorm;-use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig;-use burn::nn::conv::ConvTranspose2d;-use burn::nn::conv::ConvTranspose2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule2 { constant213: burn::module::Param<Tensor<1, Int>>,@@ -1,13 +1,12 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig;-use burn::nn::conv::ConvTranspose2d;-use burn::nn::conv::ConvTranspose2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule3 { linear5: Linear,@@ -1,5 +1,5 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-use super::*;+use super::{Submodule3, Tensor}; impl Submodule3 { pub fn forward( &self,@@ -1,5 +1,4 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear;@@ -8,6 +7,8 @@ use burn::nn::conv::ConvTranspose2d; use burn::nn::conv::ConvTranspose2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule4 { layernormalization5: LayerNorm,@@ -1,5 +1,5 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-use super::*;+use super::{Submodule4, Tensor, s}; impl Submodule4 { pub fn forward( &self,@@ -1,13 +1,12 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)]-use burn::nn::LayerNorm;-use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig; use burn::nn::conv::ConvTranspose2d; use burn::nn::conv::ConvTranspose2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule5 { constant483: burn::module::Param<Tensor<1>>,@@ -1,16 +1,8 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)]-use burn::nn::LayerNorm;-use burn::nn::LayerNormConfig;-use burn::nn::Linear;-use burn::nn::LinearConfig;-use burn::nn::PaddingConfig2d;-use burn::nn::conv::Conv2d;-use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ mod part_01; use part_01::Submodule1; mod part_02;@@ -1,5 +1,4 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear;@@ -7,10 +6,10 @@ use burn::nn::LinearConfig; use burn::nn::PaddingConfig2d; use burn::nn::conv::Conv2d; use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule1 { conv2d1: Conv2d,@@ -1,5 +1,7 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-use super::*;+use super::super::super::broadcast_shape;+use super::{Int, Submodule1, Tensor};+ impl Submodule1 { pub fn forward(&self, image: Tensor<4>) -> Tensor<4> { let conv2d1_out1 = self.conv2d1.forward(image);@@ -13,17 +15,9 @@ impl Submodule1 { let constantofshape1_out1: [i64; 3usize] = [1i64, 1i64, 1i64]; let constant3_out1 = self.constant3.val(); let expand1_out1 = {- let onnx_shape: [i64; 3usize] = constantofshape1_out1;+ let onnx_shape = constantofshape1_out1; let input_dims = constant3_out1.dims();- let mut shape = onnx_shape;- #[allow(clippy::needless_range_loop)]- for i in 0..3usize {- let dim_offset = i;- if shape[dim_offset] == 1 && input_dims[i] > 1 {- shape[dim_offset] = input_dims[i] as i64;- }- }- constant3_out1.expand(shape)+ constant3_out1.expand(broadcast_shape(onnx_shape, input_dims, 0)) }; let tile1_out1 = expand1_out1.repeat(&[32, 32, 1]); let add1_out1 = unsqueeze1_out1.add((tile1_out1).unsqueeze_dims(&[0isize]));@@ -1,16 +1,15 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig; use burn::nn::PaddingConfig2d;-use burn::nn::conv::Conv2d;-use burn::nn::conv::Conv2dConfig; use burn::nn::pool::MaxPool2d; use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule2 { layernormalization2: LayerNorm,@@ -1,16 +1,12 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig;-use burn::nn::PaddingConfig2d;-use burn::nn::conv::Conv2d;-use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule3 { layernormalization4: LayerNorm,@@ -1,16 +1,15 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig; use burn::nn::PaddingConfig2d;-use burn::nn::conv::Conv2d;-use burn::nn::conv::Conv2dConfig; use burn::nn::pool::MaxPool2d; use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule4 { layernormalization6: LayerNorm,@@ -1,16 +1,12 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig;-use burn::nn::PaddingConfig2d;-use burn::nn::conv::Conv2d;-use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule5 { layernormalization8: LayerNorm,@@ -1,5 +1,5 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-use super::*;+use super::{Int, Module, Submodule5, Tensor, s}; impl Submodule5 { pub fn forward(&self, add24_out1: Tensor<4>) -> Tensor<4> { let layernormalization8_out1 = {@@ -1,16 +1,12 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig;-use burn::nn::PaddingConfig2d;-use burn::nn::conv::Conv2d;-use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule6 { layernormalization10: LayerNorm,@@ -1,5 +1,5 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-use super::*;+use super::{Int, Module, Submodule6, Tensor, s}; impl Submodule6 { pub fn forward(&self, add29_out1: Tensor<4>) -> Tensor<4> { let layernormalization10_out1 = {@@ -1,16 +1,12 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig;-use burn::nn::PaddingConfig2d;-use burn::nn::conv::Conv2d;-use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule7 { layernormalization14: LayerNorm,@@ -1,5 +1,5 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-use super::*;+use super::{Int, Module, Submodule7, Tensor, s}; impl Submodule7 { pub fn forward(&self, add37_out1: Tensor<4>) -> Tensor<4> { let layernormalization14_out1 = {@@ -1,16 +1,15 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear; use burn::nn::LinearConfig; use burn::nn::PaddingConfig2d;-use burn::nn::conv::Conv2d;-use burn::nn::conv::Conv2dConfig; use burn::nn::pool::MaxPool2d; use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule8 { layernormalization18: LayerNorm,@@ -1,5 +1,5 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-use super::*;+use super::{Int, Module, Submodule8, Tensor, s}; impl Submodule8 { pub fn forward(&self, add45_out1: Tensor<4>) -> (Tensor<4>, Tensor<4>) { let layernormalization18_out1 = {@@ -1,5 +1,4 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)] use burn::nn::LayerNorm; use burn::nn::LayerNormConfig; use burn::nn::Linear;@@ -7,10 +6,10 @@ use burn::nn::LinearConfig; use burn::nn::PaddingConfig2d; use burn::nn::conv::Conv2d; use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule9 { layernormalization22: LayerNorm,@@ -1,5 +1,5 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-use super::*;+use super::{Int, Module, Submodule9, Tensor, s}; impl Submodule9 { pub fn forward(&self, add55_out1: Tensor<4>, transpose70_out1: Tensor<4>) -> Tensor<4> { let layernormalization22_out1 = {@@ -1,16 +1,8 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)]-use burn::nn::LayerNorm;-use burn::nn::LayerNormConfig;-use burn::nn::Linear;-use burn::nn::LinearConfig;-use burn::nn::PaddingConfig2d;-use burn::nn::conv::Conv2d;-use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule10 { constant977: burn::module::Param<Tensor<1>>,@@ -1,16 +1,11 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)]-use burn::nn::LayerNorm;-use burn::nn::LayerNormConfig;-use burn::nn::Linear;-use burn::nn::LinearConfig; use burn::nn::PaddingConfig2d; use burn::nn::conv::Conv2d; use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule11 { conv2d4: Conv2d,@@ -1,16 +1,11 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)]-use burn::nn::LayerNorm;-use burn::nn::LayerNormConfig;-use burn::nn::Linear;-use burn::nn::LinearConfig; use burn::nn::PaddingConfig2d; use burn::nn::conv::Conv2d; use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule12 { conv2d5: Conv2d,@@ -1,16 +1,11 @@ // @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.-#![allow(unused_imports)]-use burn::nn::LayerNorm;-use burn::nn::LayerNormConfig;-use burn::nn::Linear;-use burn::nn::LinearConfig; use burn::nn::PaddingConfig2d; use burn::nn::conv::Conv2d; use burn::nn::conv::Conv2dConfig;-use burn::nn::pool::MaxPool2d;-use burn::nn::pool::MaxPool2dConfig; use burn::prelude::*; +// @generated by tools/burn/generate_sam21.py (burn-onnx 0.22.0 plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.+ #[derive(Module, Debug)] pub struct Submodule13 { conv2d6: Conv2d,@@ -5,5 +5,30 @@ pub mod decoder; pub mod encoder; +fn broadcast_shape<const RANK: usize, const INPUT_RANK: usize>(+ mut requested: [i64; RANK],+ input_dimensions: [usize; INPUT_RANK],+ offset: usize,+) -> [i64; RANK] {+ for (index, input_dimension) in input_dimensions.into_iter().enumerate() {+ let dimension = &mut requested[offset + index];+ if *dimension == 1 && input_dimension > 1 {+ *dimension = input_dimension as i64;+ }+ }+ requested+}++#[cfg(test)]+mod tests {+ use super::broadcast_shape;++ #[test]+ fn broadcasts_only_singleton_dimensions_in_the_input_range() {+ assert_eq!(broadcast_shape([1, 3, 1, 7], [5, 2], 1), [1, 3, 2, 7]);+ assert_eq!(broadcast_shape([4, 6], [9, 3], 0), [4, 6]);+ }+}+ /// What this code was generated from; part of the saved-weights cache key. pub const ARCHITECTURE_REVISION: &str = "sam21-tiny-burn-0.22.1-onnx-0.22.0-v1";@@ -34,9 +34,21 @@ BURN_ONNX = '0.22.0' MAX_LINES = 700 # A file holding one whole submodule is split once it would pass this many lines. SPLIT_AT = 500+# The generated forward modules need only the symbols used by their submodule. Keep this list+# explicit so Clippy can catch accidental coupling to the graph-wide prelude.+FORWARD_IMPORTS = {+ ('decoder', 1): ('s', 'Tensor', 'Int', 'Module', 'Bool'),+ ('decoder', 3): ('Tensor',),+ ('decoder', 4): ('s', 'Tensor'),+ ('encoder', 1): ('Tensor', 'Int'),+ ('encoder', 5): ('s', 'Tensor', 'Module', 'Int'),+ ('encoder', 6): ('s', 'Tensor', 'Module', 'Int'),+ ('encoder', 7): ('s', 'Tensor', 'Module', 'Int'),+ ('encoder', 8): ('s', 'Tensor', 'Module', 'Int'),+ ('encoder', 9): ('s', 'Tensor', 'Module', 'Int'),+} HEADER = ('// @generated by tools/burn/generate_sam21.py (burn-onnx ' + BURN_ONNX + ' plus its patches) from the published SAM 2.1 tiny ONNX. Do not edit: rerun it.\n')-ALLOW = '#![allow(unused_imports)]\n' CLIPPY_ATTR = ' #[allow(clippy::let_and_return, clippy::approx_constant)]\n' NL = '\n' @@ -90,6 +102,31 @@ FLOAT_FILL = re.compile( r'\s+\(&self\.device, burn::tensor::DType::F32\),\n\s+\)\n\s+\.reshape\(\[((?:1(?:, )?)+)\]\)\n\s+\.expand\((concat\d+_out1)\);') WHERE_BCAST = re.compile( r'let (where\d+_out1) = (\w+)\.mask_where\((unsqueeze\d+_out1), (add\d+_out1)\);')+EXPAND_SHAPE = re.compile(+ r'let (expand\d+_out1) = \{\n'+ r'\s+let onnx_shape(?:\s*:\s*\[i64; \d+usize\])? = (.*?);\n'+ r'\s+let input_dims = (\w+)\.dims\(\);\n'+ r'\s+let mut shape = onnx_shape;\n'+ r'\s+#\[allow\(clippy::needless_range_loop\)\]\n'+ r'\s+for i in 0\.\.\d+usize \{\n'+ r'\s+let dim_offset = (i|\d+usize \+ i);\n'+ r'\s+if shape\[dim_offset\] == 1 && input_dims\[i\] > 1 \{\n'+ r'\s+shape\[dim_offset\] = input_dims\[i\] as i64;\n'+ r'\s+\}\n'+ r'\s+\}\n'+ r'\s+(\w+)\.expand\(shape\)\n'+ r'\s+\};', re.DOTALL)+++def simplify_broadcast_shape(m):+ """Delegate repeated ONNX broadcast-shape expansion to a small, reviewable helper."""+ output, requested, dimensions, offset, tensor = m.groups()+ offset = '0' if offset == 'i' else offset.split('usize', 1)[0]+ return (f'let {output} = {{\n'+ f' let onnx_shape = {requested};\n'+ f' let input_dims = {dimensions}.dims();\n'+ f' {tensor}.expand(broadcast_shape(onnx_shape, input_dims, {offset}))\n'+ ' };') def once(text, old, new, what):@@ -177,7 +214,20 @@ def uses_of(text): keep = [l for l in top if l.startswith("use ") and "burn_store" not in l and "::Bytes" not in l] if not keep: fail("no use lines in the generated header")- return NL.join(keep) + NL+ return keep+++def imports_for(text, imports):+ """Keep only names referenced by this item; the Burn prelude remains a prelude import."""+ used = []+ for line in imports:+ if line == 'use burn::prelude::*;':+ used.append(line)+ continue+ name = line.removesuffix(';').rsplit('::', 1)[-1]+ if re.search(r'\b' + re.escape(name) + r'\b', text):+ used.append(line)+ return NL.join(used) + NL def tidy(lines):@@ -187,7 +237,7 @@ def tidy(lines): return lines -def sub_files(number, chunk, uses):+def sub_files(graph, number, chunk, uses): """The file(s) of one submodule: whole when small, else `new` here and `forward` in a child.""" chunk = tidy(chunk) name = f'part_{number:02d}'@@ -199,13 +249,19 @@ def sub_files(number, chunk, uses): fwd -= 1 ident = re.match(r'pub struct (Submodule\d+)', chunk[1]).group(1) if len(chunk) + 8 <= SPLIT_AT:- return {f'{name}.rs': ALLOW + uses + '\n' + '\n'.join(chunk) + '\n'}+ content = '\n'.join(chunk)+ return {f'{name}.rs': imports_for(content, uses) + '\n' + content + '\n'} # Fields stay private: a child module of the struct's module can still read them. head = ('\n'.join(chunk[:struct_end + 1]) + '\n\n' + '\n'.join(chunk[impl_start:fwd]).rstrip() + '\n}\n\nmod forward;\n')- child = f'impl {ident} {{\n' + '\n'.join(chunk[fwd:impl_end]) + '\n}\n'- return {f'{name}.rs': ALLOW + uses + '\n' + head,- f'{name}/forward.rs': 'use super::*;\n' + child}+ imports = ', '.join((ident, *FORWARD_IMPORTS[(graph, number)]))+ body = '\n'.join(chunk[fwd:impl_end])+ body, replaced = EXPAND_SHAPE.subn(simplify_broadcast_shape, body)+ shape_import = 'use super::super::super::broadcast_shape;\n' if replaced else ''+ child = (f'use super::{{{imports}}};\n' + shape_import ++ f'impl {ident} {{\n' + body + '\n}\n')+ return {f'{name}.rs': imports_for(head, uses) + '\n' + head,+ f'{name}/forward.rs': child} def model_file(count, model, uses):@@ -218,7 +274,8 @@ def model_file(count, model, uses): new_impl = next(i for i, line in enumerate(model) if line == 'impl Model {' and any('pub fn new' in x for x in model[i:i + 4])) mods = ''.join(f'mod part_{n:02d};\nuse part_{n:02d}::Submodule{n};\n' for n in range(1, count + 1))- return (ALLOW + uses + '\n' + mods + '\n' + '\n'.join(model[:struct_end + 1]) + '\n\n' ++ content = '\n'.join(model[:struct_end + 1]) + '\n\n' + '\n'.join(model[new_impl:])+ return (imports_for(content, uses) + '\n' + mods + '\n' + '\n'.join(model[:struct_end + 1]) + '\n\n' + '\n'.join(model[new_impl:]) + '\n') @@ -228,7 +285,7 @@ def render(name, text): uses = uses_of(text) files = {} for n, chunk in enumerate(subs, 1):- for path, body in sub_files(n, chunk, uses).items():+ for path, body in sub_files(name, n, chunk, uses).items(): files[f'{name}/{path}'] = body files[f'{name}/mod.rs'] = model_file(len(subs), model, uses) # The clippy allowance on every `forward` is made once, on the module declaration.@@ -241,6 +298,28 @@ def top_file(): '//! The compiled SAM 2.1 tiny architecture. [`ARCHITECTURE_REVISION`] must equal the\n' '//! revision in the saved-weights cache key of `craft-ai-burn-setup`.\n\n' 'pub mod decoder;\npub mod encoder;\n\n'+ 'fn broadcast_shape<const RANK: usize, const INPUT_RANK: usize>(\n'+ ' mut requested: [i64; RANK],\n'+ ' input_dimensions: [usize; INPUT_RANK],\n'+ ' offset: usize,\n'+ ') -> [i64; RANK] {\n'+ ' for (index, input_dimension) in input_dimensions.into_iter().enumerate() {\n'+ ' let dimension = &mut requested[offset + index];\n'+ ' if *dimension == 1 && input_dimension > 1 {\n'+ ' *dimension = input_dimension as i64;\n'+ ' }\n'+ ' }\n'+ ' requested\n'+ '}\n\n'+ '#[cfg(test)]\n'+ 'mod tests {\n'+ ' use super::broadcast_shape;\n\n'+ ' #[test]\n'+ ' fn broadcasts_only_singleton_dimensions_in_the_input_range() {\n'+ ' assert_eq!(broadcast_shape([1, 3, 1, 7], [5, 2], 1), [1, 3, 2, 7]);\n'+ ' assert_eq!(broadcast_shape([4, 6], [9, 3], 0), [4, 6]);\n'+ ' }\n'+ '}\n\n' '/// What this code was generated from; part of the saved-weights cache key.\n' f'pub const ARCHITECTURE_REVISION: &str = "{REVISION}";\n')