#116

Committed 14 Aug 2025 09:33PM UTC coverage: 82.756% (+0.5%) from 82.236%

Build # #116

Build Type

push

Committed by

Commit Message

refactor: replace `run_every_nth` with `run_periodically` throughout codebase

- Rename macro for periodic execution across all modules for clearer naming.
- Introduce `unpack_shape_contract!` and update related methods to improve usability.
- Update `README` and `Cargo.toml` to reflect changes.
- Bump crate versions to `0.3.0` in alignment with new API updates.

Coverage Stats

87 of 114 new or added lines in 8 files covered. (76.32%)

2054 of 2482 relevant lines covered (82.76%)

80.55 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

93.83

/crates/bimm/src/layers/patching/patch_embed.rs

use bimm_contracts::assert_shape_contract_periodically;
use burn::config::Config;
use burn::module::Module;
use burn::nn::conv::{Conv2d, Conv2dConfig};
use burn::nn::{LayerNorm, LayerNormConfig};
use burn::prelude::{Backend, Tensor};

/// Common introspection interface for `PatchEmbed` modules.
pub trait PatchEmbedMeta {
    /// Input resolution (height, width).
    fn input_resolution(&self) -> [usize; 2];

    /// Input height.
    fn input_height(&self) -> usize {
        self.input_resolution()[0]
    }

    /// Input width.
    fn input_width(&self) -> usize {
        self.input_resolution()[1]
    }

    /// Input feature dimension size.
    fn d_input(&self) -> usize;

    /// The size of each patch.
    fn patch_size(&self) -> usize;

    /// Image resolution, measured in patches.
    fn patches_resolution(&self) -> [usize; 2] {
        [self.patches_height(), self.patches_width()]
    }

    /// Height of the image, measured in patches.
    fn patches_height(&self) -> usize {
        self.input_height() / self.patch_size()
    }

    /// Width of the image, measured in patches.
    fn patches_width(&self) -> usize {
        self.input_width() / self.patch_size()
    }

    /// Total number of patches.
    fn num_patches(&self) -> usize {
        self.patches_height() * self.patches_width()
    }

    /// Output feature dimension size.
    fn d_output(&self) -> usize;

    /// Enable patch normalization.
    fn enable_patch_norm(&self) -> bool;
}

/// Configuration for `PatchEmbed`.
#[derive(Config, Debug, Copy)]
pub struct PatchEmbedConfig {
    /// Input resolution (height, width).
    input_resolution: [usize; 2],

    /// Patch size.
    patch_size: usize,

    /// Input feature dimension size.
    d_input: usize,

    /// Output feature dimension size.
    d_output: usize,

    /// Enable patch normalization.
    #[config(default = true)]
    enable_patch_norm: bool,
}

impl PatchEmbedMeta for PatchEmbedConfig {
    fn input_resolution(&self) -> [usize; 2] {
        self.input_resolution
    }

    fn patch_size(&self) -> usize {
        self.patch_size
    }

    fn d_input(&self) -> usize {
        self.d_input
    }

    fn d_output(&self) -> usize {
        self.d_output
    }

    fn enable_patch_norm(&self) -> bool {
        self.enable_patch_norm
    }
}

impl PatchEmbedConfig {
    /// Initialize a `PatchEmbed` module with the given configuration.
    ///
    /// ## Arguments
    ///
    /// * `device` - The device on which the module will be initialized.
    ///
    /// ## Returns
    ///
    /// * A `PatchEmbed` module configured with the specified parameters.
    #[must_use]
    pub fn init<B: Backend>(
        &self,
        device: &B::Device,
    ) -> PatchEmbed<B> {
        let [h, w] = self.input_resolution;
        assert!(
            h % self.patch_size == 0 && w % self.patch_size == 0,
            "Input resolution must be divisible by patch size: {:?}",
            self.input_resolution
        );

        let stride = [self.patch_size, self.patch_size];

        PatchEmbed {
            input_resolution: self.input_resolution,
            patch_size: self.patch_size,
            projection: Conv2dConfig::new([self.d_input, self.d_output], stride)
                .with_stride(stride)
                .init(device),
            norm: match self.enable_patch_norm {
                true => Some(LayerNormConfig::new(self.d_output()).init(device)),
                false => None,
            },
        }
    }
}

/// SWIN-Transformer v2 `PatchEmbed` module.
#[derive(Module, Debug)]
pub struct PatchEmbed<B: Backend> {
    /// Input resolution (height, width).
    pub input_resolution: [usize; 2],

    /// Size of each patch.
    pub patch_size: usize,

    /// Convolutional layer for patch projection.
    pub projection: Conv2d<B>,

    /// Patch normalization layer, if enabled.
    pub norm: Option<LayerNorm<B>>,
}

impl<B: Backend> PatchEmbedMeta for PatchEmbed<B> {
    fn input_resolution(&self) -> [usize; 2] {
        self.input_resolution
    }

    fn patch_size(&self) -> usize {
        self.patch_size
    }

    fn d_input(&self) -> usize {
        self.projection.weight.dims()[1]
    }

    fn d_output(&self) -> usize {
        self.projection.weight.dims()[0]
    }

    fn enable_patch_norm(&self) -> bool {
        self.norm.is_some()
    }
}

impl<B: Backend> PatchEmbed<B> {
    /// Apply the `PatchEmbed` module to an input tensor.
    ///
    /// ## Arguments
    ///
    /// * `x` - Input tensor of shape ``(B, C, H, W)``.
    ///
    /// ## Returns
    ///
    /// * Output tensor of shape ``(B, H/patch_size * W/patch_size, d_output)``.
    #[must_use]
    pub fn forward(
        &self,
        x: Tensor<B, 4>,
    ) -> Tensor<B, 3> {
        assert_shape_contract_periodically!(
            ["batch", "d_input", "height", "width"],
            &x,
            &[
                ("d_input", self.d_input()),
                ("height", self.input_height()),
                ("width", self.input_width()),
            ]
        );

        let batch = x.dims()[0];

        let x = self.projection.forward(x);
        assert_shape_contract_periodically!(
            ["batch", "d_output", "patches_height", "patches_width"],
            &x,
            &[
                ("batch", batch),
                ("d_output", self.d_output()),
                ("patches_height", self.patches_height()),
                ("patches_width", self.patches_width()),
            ],
        );

        let x = x.flatten(2, 3);
        let x = x.swap_dims(1, 2);
        assert_shape_contract_periodically!(
            ["batch", "num_patches", "d_output"],
            &x,
            &[
                ("batch", batch),
                ("num_patches", self.num_patches()),
                ("d_output", self.d_output()),
            ],
        );

        let x = match self.norm {
            None => x,
            Some(ref norm) => norm.forward(x),
        };
        assert_shape_contract_periodically!(
            ["batch", "num_patches", "d_output"],
            &x,
            &[
                ("batch", batch),
                ("num_patches", self.num_patches()),
                ("d_output", self.d_output()),
            ],
        );

        x
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use burn::backend::NdArray;
    use burn::tensor::TensorData;

    #[test]
    fn test_patch_embed_meta() {
        let config = PatchEmbedConfig {
            input_resolution: [224, 224],
            patch_size: 16,
            d_input: 3,
            d_output: 768,
            enable_patch_norm: true,
        };

        assert_eq!(config.input_resolution(), [224, 224]);
        assert_eq!(config.patch_size(), 16);
        assert_eq!(config.d_input(), 3);
        assert_eq!(config.d_output(), 768);
        assert!(config.enable_patch_norm());
        assert_eq!(config.patches_resolution(), [14, 14]);
        assert_eq!(config.patches_height(), 14);
        assert_eq!(config.patches_width(), 14);
        assert_eq!(config.num_patches(), 196);
        assert_eq!(config.d_output(), 768);
        assert!(config.enable_patch_norm());

        let device = Default::default();
        let patch_embed = config.init::<NdArray>(&device);

        assert_eq!(patch_embed.input_resolution(), [224, 224]);
        assert_eq!(patch_embed.patch_size(), 16);
        assert_eq!(patch_embed.d_input(), 3);
        assert_eq!(patch_embed.d_output(), 768);
        assert!(patch_embed.enable_patch_norm());
        assert_eq!(patch_embed.patches_resolution(), [14, 14]);
        assert_eq!(patch_embed.patches_height(), 14);
        assert_eq!(patch_embed.patches_width(), 14);
        assert_eq!(patch_embed.num_patches(), 196);
        assert_eq!(patch_embed.d_output(), 768);
        assert!(patch_embed.enable_patch_norm());
    }

    #[should_panic(expected = "Input resolution must be divisible by patch size")]
    #[test]
    fn test_patch_embed_invalid_resolution() {
        let config = PatchEmbedConfig {
            input_resolution: [224, 223], // Invalid resolution
            patch_size: 16,
            d_input: 3,
            d_output: 768,
            enable_patch_norm: true,
        };
        let device = Default::default();
        let _d = config.init::<NdArray>(&device);
    }

    #[test]
    fn test_patch_embed_forward() {
        let config = PatchEmbedConfig {
            input_resolution: [224, 224],
            patch_size: 16,
            d_input: 3,
            d_output: 768,
            enable_patch_norm: true,
        };
        let device = Default::default();
        let patch_embed = config.init::<NdArray>(&device);

        let input = Tensor::<NdArray, 4>::from_data(
            TensorData::new(vec![1.0; 3 * 224 * 224], [1, 3, 224, 224]),
            &device,
        );

        let output = patch_embed.forward(input);
        assert_eq!(output.dims(), [1, 196, 768]);
    }

    #[test]
    fn test_patch_embed_without_norm() {
        let config = PatchEmbedConfig {
            input_resolution: [224, 224],
            patch_size: 16,
            d_input: 3,
            d_output: 768,
            enable_patch_norm: false,
        };
        let device = Default::default();
        let patch_embed = config.init::<NdArray>(&device);

        let input = Tensor::<NdArray, 4>::from_data(
            TensorData::new(vec![1.0; 3 * 224 * 224], [1, 3, 224, 224]),
            &device,
        );

        let output = patch_embed.forward(input);
        assert_eq!(output.dims(), [1, 196, 768]);
    }
}

1	use bimm_contracts::assert_shape_contract_periodically;
2	use burn::config::Config;
3	use burn::module::Module;
4	use burn::nn::conv::{Conv2d, Conv2dConfig};
5	use burn::nn::{LayerNorm, LayerNormConfig};
6	use burn::prelude::{Backend, Tensor};
7
8	/// Common introspection interface for `PatchEmbed` modules.
9	pub trait PatchEmbedMeta {
10	/// Input resolution (height, width).
11	fn input_resolution(&self) -> [usize; 2];
12
13	/// Input height.
14	fn input_height(&self) -> usize {	53✔
15	self.input_resolution()[0]	106✔
16	}
17
18	/// Input width.
19	fn input_width(&self) -> usize {	53✔
20	self.input_resolution()[1]	106✔
21	}
22
23	/// Input feature dimension size.
24	fn d_input(&self) -> usize;
25
26	/// The size of each patch.
27	fn patch_size(&self) -> usize;
28
29	/// Image resolution, measured in patches.
30	fn patches_resolution(&self) -> [usize; 2] {	4✔
31	[self.patches_height(), self.patches_width()]	12✔
32	}
33
34	/// Height of the image, measured in patches.
35	fn patches_height(&self) -> usize {	45✔
36	self.input_height() / self.patch_size()	135✔
37	}
38
39	/// Width of the image, measured in patches.
40	fn patches_width(&self) -> usize {	45✔
41	self.input_width() / self.patch_size()	135✔
42	}
43
44	/// Total number of patches.
45	fn num_patches(&self) -> usize {	22✔
46	self.patches_height() * self.patches_width()	66✔
47	}
48
49	/// Output feature dimension size.
50	fn d_output(&self) -> usize;
51
52	/// Enable patch normalization.
53	fn enable_patch_norm(&self) -> bool;
54	}
55
56	/// Configuration for `PatchEmbed`.
57	#[derive(Config, Debug, Copy)]
58	pub struct PatchEmbedConfig {
59	/// Input resolution (height, width).
60	input_resolution: [usize; 2],
61
62	/// Patch size.
63	patch_size: usize,
64
65	/// Input feature dimension size.
66	d_input: usize,
67
68	/// Output feature dimension size.
69	d_output: usize,
70
71	/// Enable patch normalization.
72	#[config(default = true)]
73	enable_patch_norm: bool,
74	}
75
76	impl PatchEmbedMeta for PatchEmbedConfig {
77	fn input_resolution(&self) -> [usize; 2] {	29✔
78	self.input_resolution	29✔
79	}
80
81	fn patch_size(&self) -> usize {	26✔
82	self.patch_size	26✔
83	}
84
85	fn d_input(&self) -> usize {	2✔
86	self.d_input	2✔
87	}
88
89	fn d_output(&self) -> usize {	16✔
90	self.d_output	16✔
91	}
92
93	fn enable_patch_norm(&self) -> bool {	3✔
94	self.enable_patch_norm	3✔
95	}
96	}
97
98	impl PatchEmbedConfig {
99	/// Initialize a `PatchEmbed` module with the given configuration.
100	///
101	/// ## Arguments
102	///
103	/// * `device` - The device on which the module will be initialized.
104	///
105	/// ## Returns
106	///
107	/// * A `PatchEmbed` module configured with the specified parameters.
108	#[must_use]
109	pub fn init<B: Backend>(	10✔
110	&self,
111	device: &B::Device,
112	) -> PatchEmbed<B> {
113	let [h, w] = self.input_resolution;	30✔
114	assert!(	10✔
115	h % self.patch_size == 0 && w % self.patch_size == 0,	20✔
116	"Input resolution must be divisible by patch size: {:?}",	1✔
117	self.input_resolution	×
118	);
119
120	let stride = [self.patch_size, self.patch_size];	18✔
121
122	PatchEmbed {
123	input_resolution: self.input_resolution,	9✔
124	patch_size: self.patch_size,	9✔
125	projection: Conv2dConfig::new([self.d_input, self.d_output], stride)	36✔
126	.with_stride(stride)
127	.init(device),
128	norm: match self.enable_patch_norm {	9✔
129	true => Some(LayerNormConfig::new(self.d_output()).init(device)),
130	false => None,
131	},
132	}
133	}
134	}
135
136	/// SWIN-Transformer v2 `PatchEmbed` module.
137	#[derive(Module, Debug)]
138	pub struct PatchEmbed<B: Backend> {
139	/// Input resolution (height, width).
140	pub input_resolution: [usize; 2],
141
142	/// Size of each patch.
143	pub patch_size: usize,
144
145	/// Convolutional layer for patch projection.
146	pub projection: Conv2d<B>,
147
148	/// Patch normalization layer, if enabled.
149	pub norm: Option<LayerNorm<B>>,
150	}
151
152	impl<B: Backend> PatchEmbedMeta for PatchEmbed<B> {
153	fn input_resolution(&self) -> [usize; 2] {	91✔
154	self.input_resolution	91✔
155	}
156
157	fn patch_size(&self) -> usize {	69✔
158	self.patch_size	69✔
159	}
160
161	fn d_input(&self) -> usize {	13✔
162	self.projection.weight.dims()[1]	13✔
163	}
164
165	fn d_output(&self) -> usize {	30✔
166	self.projection.weight.dims()[0]	30✔
167	}
168
169	fn enable_patch_norm(&self) -> bool {	4✔
170	self.norm.is_some()	8✔
171	}
172	}
173
174	impl<B: Backend> PatchEmbed<B> {
175	/// Apply the `PatchEmbed` module to an input tensor.
176	///
177	/// ## Arguments
178	///
179	/// * `x` - Input tensor of shape ``(B, C, H, W)``.
180	///
181	/// ## Returns
182	///
183	/// * Output tensor of shape ``(B, H/patch_size * W/patch_size, d_output)``.
184	#[must_use]
185	pub fn forward(	8✔
186	&self,
187	x: Tensor<B, 4>,
188	) -> Tensor<B, 3> {
189	assert_shape_contract_periodically!(	8✔
NEW 190	["batch", "d_input", "height", "width"],	×
191	&x,	8✔
192	&[	8✔
193	("d_input", self.d_input()),	24✔
194	("height", self.input_height()),	24✔
195	("width", self.input_width()),	16✔
196	]
197	);
198
199	let batch = x.dims()[0];	24✔
200
201	let x = self.projection.forward(x);	32✔
202	assert_shape_contract_periodically!(	8✔
NEW 203	["batch", "d_output", "patches_height", "patches_width"],	×
204	&x,	8✔
205	&[	8✔
206	("batch", batch),	16✔
207	("d_output", self.d_output()),	24✔
208	("patches_height", self.patches_height()),	24✔
209	("patches_width", self.patches_width()),	16✔
210	],
211	);
212
213	let x = x.flatten(2, 3);	24✔
214	let x = x.swap_dims(1, 2);	24✔
215	assert_shape_contract_periodically!(	8✔
NEW 216	["batch", "num_patches", "d_output"],	×
217	&x,	8✔
218	&[	8✔
219	("batch", batch),	16✔
220	("num_patches", self.num_patches()),	24✔
221	("d_output", self.d_output()),	16✔
222	],
223	);
224
225	let x = match self.norm {	16✔
226	None => x,	2✔
227	Some(ref norm) => norm.forward(x),	24✔
228	};
229	assert_shape_contract_periodically!(	8✔
NEW 230	["batch", "num_patches", "d_output"],	×
231	&x,	8✔
232	&[	8✔
233	("batch", batch),	16✔
234	("num_patches", self.num_patches()),	24✔
235	("d_output", self.d_output()),	16✔
236	],
237	);
238
239	x	8✔
240	}
241	}
242
243	#[cfg(test)]
244	mod tests {
245	use super::*;
246	use burn::backend::NdArray;
247	use burn::tensor::TensorData;
248
249	#[test]
250	fn test_patch_embed_meta() {
251	let config = PatchEmbedConfig {
252	input_resolution: [224, 224],
253	patch_size: 16,
254	d_input: 3,
255	d_output: 768,
256	enable_patch_norm: true,
257	};
258
259	assert_eq!(config.input_resolution(), [224, 224]);
260	assert_eq!(config.patch_size(), 16);
261	assert_eq!(config.d_input(), 3);
262	assert_eq!(config.d_output(), 768);
263	assert!(config.enable_patch_norm());
264	assert_eq!(config.patches_resolution(), [14, 14]);
265	assert_eq!(config.patches_height(), 14);
266	assert_eq!(config.patches_width(), 14);
267	assert_eq!(config.num_patches(), 196);
268	assert_eq!(config.d_output(), 768);
269	assert!(config.enable_patch_norm());
270
271	let device = Default::default();
272	let patch_embed = config.init::<NdArray>(&device);
273
274	assert_eq!(patch_embed.input_resolution(), [224, 224]);
275	assert_eq!(patch_embed.patch_size(), 16);
276	assert_eq!(patch_embed.d_input(), 3);
277	assert_eq!(patch_embed.d_output(), 768);
278	assert!(patch_embed.enable_patch_norm());
279	assert_eq!(patch_embed.patches_resolution(), [14, 14]);
280	assert_eq!(patch_embed.patches_height(), 14);
281	assert_eq!(patch_embed.patches_width(), 14);
282	assert_eq!(patch_embed.num_patches(), 196);
283	assert_eq!(patch_embed.d_output(), 768);
284	assert!(patch_embed.enable_patch_norm());
285	}
286
287	#[should_panic(expected = "Input resolution must be divisible by patch size")]
288	#[test]
289	fn test_patch_embed_invalid_resolution() {
290	let config = PatchEmbedConfig {
291	input_resolution: [224, 223], // Invalid resolution
292	patch_size: 16,
293	d_input: 3,
294	d_output: 768,
295	enable_patch_norm: true,
296	};
297	let device = Default::default();
298	let _d = config.init::<NdArray>(&device);
299	}
300
301	#[test]
302	fn test_patch_embed_forward() {
303	let config = PatchEmbedConfig {
304	input_resolution: [224, 224],
305	patch_size: 16,
306	d_input: 3,
307	d_output: 768,
308	enable_patch_norm: true,
309	};
310	let device = Default::default();
311	let patch_embed = config.init::<NdArray>(&device);
312
313	let input = Tensor::<NdArray, 4>::from_data(
314	TensorData::new(vec![1.0; 3 * 224 * 224], [1, 3, 224, 224]),
315	&device,
316	);
317
318	let output = patch_embed.forward(input);
319	assert_eq!(output.dims(), [1, 196, 768]);
320	}
321
322	#[test]
323	fn test_patch_embed_without_norm() {
324	let config = PatchEmbedConfig {
325	input_resolution: [224, 224],
326	patch_size: 16,
327	d_input: 3,
328	d_output: 768,
329	enable_patch_norm: false,
330	};
331	let device = Default::default();
332	let patch_embed = config.init::<NdArray>(&device);
333
334	let input = Tensor::<NdArray, 4>::from_data(
335	TensorData::new(vec![1.0; 3 * 224 * 224], [1, 3, 224, 224]),
336	&device,
337	);
338
339	let output = patch_embed.forward(input);
340	assert_eq!(output.dims(), [1, 196, 768]);
341	}
342	}

crutcher / bimm / #116

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous