>,
diff --git a/crates/bimm-firehose/src/burn/path_scanning.rs b/crates/bimm-firehose/src/burn/path_scanning.rs
index bab900e6..dd54ed05 100644
--- a/crates/bimm-firehose/src/burn/path_scanning.rs
+++ b/crates/bimm-firehose/src/burn/path_scanning.rs
@@ -1,6 +1,12 @@
-use burn::data::dataset::InMemDataset;
-use burn::data::dataset::vision::ImageLoaderError;
-use std::collections::{HashMap, HashSet};
+use std::collections::{
+ HashMap,
+ HashSet,
+};
+
+use burn::data::dataset::{
+ InMemDataset,
+ vision::ImageLoaderError,
+};
/// Scan a folder of ``$ROOT/$CLASS/$IMG.{jpg,png}`` into an `InMemDataset`.
pub fn image_dataset_for_folder(root: P) -> anyhow::Result>
diff --git a/crates/bimm-firehose/src/core/mod.rs b/crates/bimm-firehose/src/core/mod.rs
index 035750b6..61a7a080 100644
--- a/crates/bimm-firehose/src/core/mod.rs
+++ b/crates/bimm-firehose/src/core/mod.rs
@@ -12,8 +12,12 @@ pub mod values;
// TODO: Work out what the `$crate::core::*` re-exports should be.
pub use rows::{
- FirehoseBatchTransaction, FirehoseRow, FirehoseRowBatch, FirehoseRowReader,
- FirehoseRowTransaction, FirehoseRowWriter,
+ FirehoseBatchTransaction,
+ FirehoseRow,
+ FirehoseRowBatch,
+ FirehoseRowReader,
+ FirehoseRowTransaction,
+ FirehoseRowWriter,
};
pub use schema::FirehoseTableSchema;
pub use values::FirehoseValue;
diff --git a/crates/bimm-firehose/src/core/operations/environment.rs b/crates/bimm-firehose/src/core/operations/environment.rs
index 6273a647..7c32a66b 100644
--- a/crates/bimm-firehose/src/core/operations/environment.rs
+++ b/crates/bimm-firehose/src/core/operations/environment.rs
@@ -1,14 +1,33 @@
-use crate::core::operations::factory::{FirehoseOperatorFactory, FirehoseOperatorInitContext};
-use crate::core::operations::operator::FirehoseOperator;
-use crate::core::operations::planner::OperationPlan;
-use crate::core::operations::signature::FirehoseOperatorSignature;
-use crate::core::schema::{BuildPlan, DataTypeDescription, FirehoseTableSchema};
-use anyhow::{Context, bail};
-use std::collections::BTreeMap;
-use std::fmt::Debug;
-use std::sync::Arc;
-
-/// `OpEnvironment` is a trait that provides access to a collection of operator bindings.
+use std::{
+ collections::BTreeMap,
+ fmt::Debug,
+ sync::Arc,
+};
+
+use anyhow::{
+ Context,
+ bail,
+};
+
+use crate::core::{
+ operations::{
+ factory::{
+ FirehoseOperatorFactory,
+ FirehoseOperatorInitContext,
+ },
+ operator::FirehoseOperator,
+ planner::OperationPlan,
+ signature::FirehoseOperatorSignature,
+ },
+ schema::{
+ BuildPlan,
+ DataTypeDescription,
+ FirehoseTableSchema,
+ },
+};
+
+/// `OpEnvironment` is a trait that provides access to a collection of operator
+/// bindings.
pub trait FirehoseOperatorEnvironment: Debug + Send + Sync {
/// Returns a reference to the map of operator bindings.
// TODO: This should be an iterator.
@@ -24,7 +43,8 @@ pub trait FirehoseOperatorEnvironment: Debug + Send + Sync {
///
/// # Returns
///
- /// An `anyhow::Result>` containing the operator factory.
+ /// An `anyhow::Result>` containing the
+ /// operator factory.
fn lookup_operator_factory(
&self,
operator_id: &str,
@@ -43,7 +63,8 @@ pub trait FirehoseOperatorEnvironment: Debug + Send + Sync {
///
/// # Arguments
///
- /// * `context` - The context containing the build plan and input/output types.
+ /// * `context` - The context containing the build plan and input/output
+ /// types.
///
/// # Returns
///
@@ -59,11 +80,13 @@ pub trait FirehoseOperatorEnvironment: Debug + Send + Sync {
///
/// # Arguments
///
- /// * `plan_context` - The context containing the build plan and input/output types.
+ /// * `plan_context` - The context containing the build plan and
+ /// input/output types.
///
/// # Returns
///
- /// An `anyhow::Result>` containing the initialized operator.
+ /// An `anyhow::Result>` containing the
+ /// initialized operator.
fn init_operator(
&self,
plan_context: BuildPlanContext,
@@ -84,11 +107,13 @@ pub trait FirehoseOperatorEnvironment: Debug + Send + Sync {
/// # Arguments
///
/// * `schema` - A mutable reference to the `TableSchema` to be extended.
- /// * `planner` - An `OperationPlanner` that contains the details of the operation to be planned.
+ /// * `planner` - An `OperationPlanner` that contains the details of the
+ /// operation to be planned.
///
/// # Returns
///
- /// An `anyhow::Result` containing the build plan for the operation.
+ /// An `anyhow::Result` containing the build plan for the
+ /// operation.
fn apply_plan_to_schema(
&self,
schema: &mut FirehoseTableSchema,
@@ -115,7 +140,8 @@ pub trait FirehoseOperatorEnvironment: Debug + Send + Sync {
}
}
-/// `MapOpEnvironment` is a simple implementation of `OpEnvironment` that uses a `BTreeMap` to store operators.
+/// `MapOpEnvironment` is a simple implementation of `OpEnvironment` that uses a
+/// `BTreeMap` to store operators.
#[derive(Debug)]
pub struct MapOpEnvironment {
/// A map of operator IDs to their corresponding operator factories.
@@ -161,7 +187,8 @@ impl MapOpEnvironment {
///
/// # Returns
///
- /// An `anyhow::Result<()>` indicating success or containing an error if the binding already exists.
+ /// An `anyhow::Result<()>` indicating success or containing an error if the
+ /// binding already exists.
pub fn add_operator(
&mut self,
factory: Arc,
@@ -182,7 +209,8 @@ impl MapOpEnvironment {
///
/// # Returns
///
- /// An `anyhow::Result<()>` indicating success or containing an error if any binding fails to be added.
+ /// An `anyhow::Result<()>` indicating success or containing an error if any
+ /// binding fails to be added.
pub fn add_all_operators(
&mut self,
factories: Vec>,
@@ -215,7 +243,8 @@ pub struct BuildPlanContext {
}
impl BuildPlanContext {
- /// Creates a new `OperationInitPlanContext` with the given table schema and build plan.
+ /// Creates a new `OperationInitPlanContext` with the given table schema and
+ /// build plan.
pub fn new(
table_schema: Arc,
build_plan: Arc,
@@ -245,7 +274,8 @@ impl BuildPlanContext {
///
/// # Returns
///
- /// An `anyhow::Result` containing the bound context.
+ /// An `anyhow::Result` containing the bound
+ /// context.
pub fn bind_signature(
self,
signature: &FirehoseOperatorSignature,
@@ -253,7 +283,8 @@ impl BuildPlanContext {
OperationInitializationContext::init(self, signature.clone())
}
- /// Computes the input types for the operator based on the build plan and table schema.
+ /// Computes the input types for the operator based on the build plan and
+ /// table schema.
pub fn input_types(&self) -> BTreeMap {
self.build_plan
.inputs
@@ -267,7 +298,8 @@ impl BuildPlanContext {
.collect()
}
- /// Computes the output types for the operator based on the build plan and table schema.
+ /// Computes the output types for the operator based on the build plan and
+ /// table schema.
pub fn output_types(&self) -> BTreeMap {
self.build_plan
.outputs
@@ -312,7 +344,8 @@ impl FirehoseOperatorInitContext for OperationInitializationContext {
}
impl OperationInitializationContext {
- /// Creates a new `OperationInitSignatureContext` with the given plan context and signature.
+ /// Creates a new `OperationInitSignatureContext` with the given plan
+ /// context and signature.
///
/// # Returns
///
diff --git a/crates/bimm-firehose/src/core/operations/executor.rs b/crates/bimm-firehose/src/core/operations/executor.rs
index 15cc2909..b04fc4ca 100644
--- a/crates/bimm-firehose/src/core/operations/executor.rs
+++ b/crates/bimm-firehose/src/core/operations/executor.rs
@@ -1,9 +1,16 @@
-use crate::core::operations::environment::FirehoseOperatorEnvironment;
-use crate::core::operations::operator::OperationRunner;
-use crate::core::rows::FirehoseRowBatch;
-use crate::core::schema::FirehoseTableSchema;
-use std::fmt::Debug;
-use std::sync::Arc;
+use std::{
+ fmt::Debug,
+ sync::Arc,
+};
+
+use crate::core::{
+ operations::{
+ environment::FirehoseOperatorEnvironment,
+ operator::OperationRunner,
+ },
+ rows::FirehoseRowBatch,
+ schema::FirehoseTableSchema,
+};
/// Trait for executing a batch of operations on a `RowBatch`.
pub trait FirehoseBatchExecutor: Debug + Send + Sync {
@@ -37,7 +44,8 @@ pub struct SequentialBatchExecutor {
}
impl SequentialBatchExecutor {
- /// Creates a new `DefaultBatchExecutor` with the given operator environment.
+ /// Creates a new `DefaultBatchExecutor` with the given operator
+ /// environment.
pub fn new(
schema: Arc,
environment: Arc,
diff --git a/crates/bimm-firehose/src/core/operations/factory.rs b/crates/bimm-firehose/src/core/operations/factory.rs
index 7ba5a69e..f5742596 100644
--- a/crates/bimm-firehose/src/core/operations/factory.rs
+++ b/crates/bimm-firehose/src/core/operations/factory.rs
@@ -1,10 +1,21 @@
-use crate::core::operations::operator::FirehoseOperator;
-use crate::core::operations::signature::FirehoseOperatorSignature;
-use crate::core::schema::{BuildPlan, FirehoseTableSchema};
+use std::{
+ fmt::Debug,
+ marker::PhantomData,
+};
+
use anyhow::Context;
use serde::de::DeserializeOwned;
-use std::fmt::Debug;
-use std::marker::PhantomData;
+
+use crate::core::{
+ operations::{
+ operator::FirehoseOperator,
+ signature::FirehoseOperatorSignature,
+ },
+ schema::{
+ BuildPlan,
+ FirehoseTableSchema,
+ },
+};
/// A factory for creating `FirehoseOperator` instances from a specification.
pub trait FirehoseOperatorFactory: Debug + Send + Sync {
@@ -19,16 +30,19 @@ pub trait FirehoseOperatorFactory: Debug + Send + Sync {
/// Returns the operator specification.
fn signature(&self) -> &FirehoseOperatorSignature;
- /// Inits a build plan against the input and output types using an `OpInitContext`.
+ /// Inits a build plan against the input and output types using an
+ /// `OpInitContext`.
///
/// # Arguments
///
- /// * `context` - The context containing the build plan and input/output types.
+ /// * `context` - The context containing the build plan and input/output
+ /// types.
///
/// # Returns
///
/// A `Result, String>` where:
- /// * `Ok` contains a boxed operator that implements the `BuildOperator` trait,
+ /// * `Ok` contains a boxed operator that implements the `BuildOperator`
+ /// trait,
/// * `Err` contains an error message if the initialization fails.
fn init(
&self,
@@ -51,7 +65,8 @@ pub trait FirehoseOperatorInitContext {
fn signature(&self) -> &FirehoseOperatorSignature;
}
-/// A simple operator factory for types implementing `DeserializeOwned` and `FirehoseOperator`.
+/// A simple operator factory for types implementing `DeserializeOwned` and
+/// `FirehoseOperator`.
#[derive(Debug)]
pub struct SimpleConfigOperatorFactory
where
@@ -68,7 +83,8 @@ impl SimpleConfigOperatorFactory
where
T: DeserializeOwned + FirehoseOperator,
{
- /// Creates a new `SpecConfigOpBinding` with the given operator specification.
+ /// Creates a new `SpecConfigOpBinding` with the given operator
+ /// specification.
pub fn new(spec: FirehoseOperatorSignature) -> Self {
if spec.operator_id.is_none() {
panic!("OperatorSpec must have an operator_id");
@@ -106,12 +122,25 @@ where
#[cfg(test)]
mod tests {
- use crate::core::operations::factory::SimpleConfigOperatorFactory;
- use crate::core::operations::operator::FirehoseOperator;
- use crate::core::operations::signature::{FirehoseOperatorSignature, ParameterSpec};
- use crate::core::rows::FirehoseRowTransaction;
- use crate::define_firehose_operator_id;
- use serde::{Deserialize, Serialize};
+ use serde::{
+ Deserialize,
+ Serialize,
+ };
+
+ use crate::{
+ core::{
+ operations::{
+ factory::SimpleConfigOperatorFactory,
+ operator::FirehoseOperator,
+ signature::{
+ FirehoseOperatorSignature,
+ ParameterSpec,
+ },
+ },
+ rows::FirehoseRowTransaction,
+ },
+ define_firehose_operator_id,
+ };
define_firehose_operator_id!(TEST_OP);
diff --git a/crates/bimm-firehose/src/core/operations/mod.rs b/crates/bimm-firehose/src/core/operations/mod.rs
index 266da069..9f106e15 100644
--- a/crates/bimm-firehose/src/core/operations/mod.rs
+++ b/crates/bimm-firehose/src/core/operations/mod.rs
@@ -17,9 +17,10 @@ pub mod signature;
///
/// # Arguments
///
-/// * `$name`: The name of the operator ID to define;
-/// will create a self-referential static string constant.
-/// * `$constructor`: A closure that returns an `Arc`.
+/// * `$name`: The name of the operator ID to define; will create a
+/// self-referential static string constant.
+/// * `$constructor`: A closure that returns an `Arc`.
///
/// This macro combines the functionality of `define_firehose_operator_id`
/// and `register_firehose_operator_factory`.
@@ -33,12 +34,12 @@ macro_rules! define_firehose_operator {
/// Define a self-referential operator ID.
///
-/// The id will be defined as a static string constant that refers to its own namespace path.
+/// The id will be defined as a static string constant that refers to its own
+/// namespace path.
///
/// # Arguments
///
/// * `$name`: The name of the operator ID to define.
-///
#[macro_export]
macro_rules! define_firehose_operator_id {
($name:ident) => {
@@ -49,7 +50,8 @@ macro_rules! define_firehose_operator_id {
/// Macro to register a default operator factory.
///
/// Builders which do not require runtime configuration can be registered
-/// using this macro; and collected globally using `list_default_operator_builders`.
+/// using this macro; and collected globally using
+/// `list_default_operator_builders`.
///
/// You can also collect a default environment with all registered builders
/// using `new_default_operator_environment`.
@@ -71,23 +73,50 @@ macro_rules! register_firehose_operator_factory {
#[cfg(test)]
mod tests {
- use crate::core::operations::environment::{FirehoseOperatorEnvironment, MapOpEnvironment};
- use crate::core::operations::factory::SimpleConfigOperatorFactory;
- use crate::core::operations::operator::OperationRunner;
- use crate::core::operations::operator::{FirehoseOperator, OperatorSchedulingMetadata};
- use crate::core::operations::signature::{FirehoseOperatorSignature, ParameterSpec};
-
- use crate::core::FirehoseValue;
- use crate::core::rows::{
- FirehoseRowBatch, FirehoseRowReader, FirehoseRowTransaction, FirehoseRowWriter,
+ use std::{
+ collections::BTreeMap,
+ fmt::Debug,
+ sync::Arc,
};
- use crate::core::schema::{BuildPlan, ColumnSchema, DataTypeDescription, FirehoseTableSchema};
+
// use crate::define_firehose_operator_id;
use indoc::indoc;
- use serde::{Deserialize, Serialize};
- use std::collections::BTreeMap;
- use std::fmt::Debug;
- use std::sync::Arc;
+ use serde::{
+ Deserialize,
+ Serialize,
+ };
+
+ use crate::core::{
+ FirehoseValue,
+ operations::{
+ environment::{
+ FirehoseOperatorEnvironment,
+ MapOpEnvironment,
+ },
+ factory::SimpleConfigOperatorFactory,
+ operator::{
+ FirehoseOperator,
+ OperationRunner,
+ OperatorSchedulingMetadata,
+ },
+ signature::{
+ FirehoseOperatorSignature,
+ ParameterSpec,
+ },
+ },
+ rows::{
+ FirehoseRowBatch,
+ FirehoseRowReader,
+ FirehoseRowTransaction,
+ FirehoseRowWriter,
+ },
+ schema::{
+ BuildPlan,
+ ColumnSchema,
+ DataTypeDescription,
+ FirehoseTableSchema,
+ },
+ };
define_firehose_operator_id!(ADD);
diff --git a/crates/bimm-firehose/src/core/operations/operator.rs b/crates/bimm-firehose/src/core/operations/operator.rs
index d79a8730..53cb208a 100644
--- a/crates/bimm-firehose/src/core/operations/operator.rs
+++ b/crates/bimm-firehose/src/core/operations/operator.rs
@@ -1,10 +1,31 @@
-use crate::core::operations::environment::{BuildPlanContext, FirehoseOperatorEnvironment};
-use crate::core::operations::signature::FirehoseOperatorSignature;
-use crate::core::rows::{FirehoseBatchTransaction, FirehoseRowBatch, FirehoseRowTransaction};
-use crate::core::schema::{BuildPlan, FirehoseTableSchema};
-use serde::{Deserialize, Serialize};
-use std::fmt::Debug;
-use std::sync::Arc;
+use std::{
+ fmt::Debug,
+ sync::Arc,
+};
+
+use serde::{
+ Deserialize,
+ Serialize,
+};
+
+use crate::core::{
+ operations::{
+ environment::{
+ BuildPlanContext,
+ FirehoseOperatorEnvironment,
+ },
+ signature::FirehoseOperatorSignature,
+ },
+ rows::{
+ FirehoseBatchTransaction,
+ FirehoseRowBatch,
+ FirehoseRowTransaction,
+ },
+ schema::{
+ BuildPlan,
+ FirehoseTableSchema,
+ },
+};
/// Scheduling metadata for an operator.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -27,7 +48,8 @@ pub trait FirehoseOperator: 'static + Send + Sync + Debug {
///
/// # Returns
///
- /// An `anyhow::Result<()>` indicating success or containing an error if the operation fails.
+ /// An `anyhow::Result<()>` indicating success or containing an error if the
+ /// operation fails.
#[must_use]
fn apply_to_batch(
&self,
@@ -44,7 +66,8 @@ pub trait FirehoseOperator: 'static + Send + Sync + Debug {
///
/// # Returns
///
- /// An `anyhow::Result<()>` indicating success or containing an error if the operation fails.
+ /// An `anyhow::Result<()>` indicating success or containing an error if the
+ /// operation fails.
#[must_use]
fn apply_to_row(
&self,
@@ -52,7 +75,8 @@ pub trait FirehoseOperator: 'static + Send + Sync + Debug {
) -> anyhow::Result<()>;
}
-/// Represents a schema + instantiated column operator for a particular build plan.
+/// Represents a schema + instantiated column operator for a particular build
+/// plan.
pub struct OperationRunner {
/// The table schema that this operator is bound to.
pub table_schema: Arc,
@@ -87,17 +111,21 @@ impl Debug for OperationRunner {
}
impl OperationRunner {
- /// Create a new `BoundPlanBuilder` by binding a `BuildPlan` to a `BimmTableSchema`.
+ /// Create a new `BoundPlanBuilder` by binding a `BuildPlan` to a
+ /// `BimmTableSchema`.
///
/// # Arguments
///
/// * `table_schema` - The schema of the table to which this plan is bound.
- /// * `build_plan` - The build plan that describes the operator and its inputs/outputs.
- /// * `env` - An environment that can create the operator based on the build plan.
+ /// * `build_plan` - The build plan that describes the operator and its
+ /// inputs/outputs.
+ /// * `env` - An environment that can create the operator based on the build
+ /// plan.
///
/// # Returns
///
- /// A result containing a `BoundPlanBuilder` if successful, or an error message if the binding fails.
+ /// A result containing a `BoundPlanBuilder` if successful, or an error
+ /// message if the binding fails.
#[must_use]
pub fn new_for_plan(
table_schema: Arc,
@@ -134,7 +162,8 @@ impl OperationRunner {
///
/// # Arguments
///
- /// * `rows` - A mutable slice of `BimmRow` instances that will be processed by the operator.
+ /// * `rows` - A mutable slice of `BimmRow` instances that will be processed
+ /// by the operator.
///
/// # Returns
///
diff --git a/crates/bimm-firehose/src/core/operations/planner.rs b/crates/bimm-firehose/src/core/operations/planner.rs
index 78041aaf..df6ecda5 100644
--- a/crates/bimm-firehose/src/core/operations/planner.rs
+++ b/crates/bimm-firehose/src/core/operations/planner.rs
@@ -1,9 +1,22 @@
-use crate::core::operations::environment::FirehoseOperatorEnvironment;
-use crate::core::operations::signature::FirehoseOperatorSignature;
-use crate::core::schema::{BuildPlan, ColumnSchema, FirehoseTableSchema};
-use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
+use serde::{
+ Deserialize,
+ Serialize,
+};
+
+use crate::core::{
+ operations::{
+ environment::FirehoseOperatorEnvironment,
+ signature::FirehoseOperatorSignature,
+ },
+ schema::{
+ BuildPlan,
+ ColumnSchema,
+ FirehoseTableSchema,
+ },
+};
+
/// A builder for constructing a call to an operator in a `BuildPlan`.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct OperationPlan {
@@ -85,7 +98,8 @@ impl OperationPlan {
///
/// # Arguments
///
- /// * `config` - The configuration to be added, which must implement `Serialize`.
+ /// * `config` - The configuration to be added, which must implement
+ /// `Serialize`.
///
/// # Returns
///
@@ -136,12 +150,15 @@ impl OperationPlan {
///
/// # Arguments
///
- /// * `schema` - The mutable reference to the table schema to which the operation will be applied.
- /// * `env` - The environment that can create the operator based on the build plan.
+ /// * `schema` - The mutable reference to the table schema to which the
+ /// operation will be applied.
+ /// * `env` - The environment that can create the operator based on the
+ /// build plan.
///
/// # Returns
///
- /// A result containing a `BuildPlan` if successful, or an error message if the operation fails.
+ /// A result containing a `BuildPlan` if successful, or an error message if
+ /// the operation fails.
pub fn apply_to_schema(
self,
schema: &mut FirehoseTableSchema,
diff --git a/crates/bimm-firehose/src/core/operations/registration.rs b/crates/bimm-firehose/src/core/operations/registration.rs
index 9ec8abc6..1664c88d 100644
--- a/crates/bimm-firehose/src/core/operations/registration.rs
+++ b/crates/bimm-firehose/src/core/operations/registration.rs
@@ -1,9 +1,13 @@
+use std::{
+ fmt::Debug,
+ sync::Arc,
+};
+
use crate::core::operations::factory::FirehoseOperatorFactory;
-use std::fmt::Debug;
-use std::sync::Arc;
// This leverages the `inventory` crate to define a collection scheme for
-// later invocations of `inventory::submit! { }` by the registration macro.
+// later invocations of `inventory::submit! { }` by the registration
+// macro.
inventory::collect!(FirehoseOperatorFactoryRegistration);
/// Struct describing a name to constructor for an operator builder.
@@ -56,12 +60,23 @@ impl FirehoseOperatorFactoryRegistration {
#[cfg(test)]
mod tests {
+ use serde::{
+ Deserialize,
+ Serialize,
+ };
+
use super::*;
- use crate::core::operations::factory::SimpleConfigOperatorFactory;
- use crate::core::operations::operator::FirehoseOperator;
- use crate::core::operations::signature::{FirehoseOperatorSignature, ParameterSpec};
- use crate::define_firehose_operator;
- use serde::{Deserialize, Serialize};
+ use crate::{
+ core::operations::{
+ factory::SimpleConfigOperatorFactory,
+ operator::FirehoseOperator,
+ signature::{
+ FirehoseOperatorSignature,
+ ParameterSpec,
+ },
+ },
+ define_firehose_operator,
+ };
define_firehose_operator!(
EXAMPLE_OPERATOR,
diff --git a/crates/bimm-firehose/src/core/operations/signature.rs b/crates/bimm-firehose/src/core/operations/signature.rs
index ff76a326..a734555b 100644
--- a/crates/bimm-firehose/src/core/operations/signature.rs
+++ b/crates/bimm-firehose/src/core/operations/signature.rs
@@ -1,8 +1,20 @@
-use crate::core::schema::{BuildPlan, ColumnSchema, DataTypeDescription};
-use anyhow::{Context, bail};
-use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
+use anyhow::{
+ Context,
+ bail,
+};
+use serde::{
+ Deserialize,
+ Serialize,
+};
+
+use crate::core::schema::{
+ BuildPlan,
+ ColumnSchema,
+ DataTypeDescription,
+};
+
/// Defines a single parameter specification
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ParameterSpec {
@@ -26,11 +38,13 @@ impl ParameterSpec {
/// # Parameters
///
/// - `name`: The name of the parameter.
- /// - `T`: The type of the parameter, which is used to determine the data type description.
+ /// - `T`: The type of the parameter, which is used to determine the data
+ /// type description.
///
/// # Returns
///
- /// A new `ParameterSpec` instance with the specified name, data type, and required arity.
+ /// A new `ParameterSpec` instance with the specified name, data type, and
+ /// required arity.
pub fn new(name: &str) -> Self {
Self {
name: name.to_string(),
@@ -101,7 +115,8 @@ impl FirehoseOperatorSignature {
}
}
- /// Creates a new `FirehoseOperatorSignature` with the specified operator ID.
+ /// Creates a new `FirehoseOperatorSignature` with the specified operator
+ /// ID.
pub fn from_operator_id(operator_id: &str) -> Self {
Self {
operator_id: Some(operator_id.into()),
@@ -137,17 +152,21 @@ impl FirehoseOperatorSignature {
}
}
- /// Internal helper to add a parameter specification to the list of inputs or outputs.
+ /// Internal helper to add a parameter specification to the list of inputs
+ /// or outputs.
///
/// # Parameters
///
/// * `spec`: The parameter specification to add.
- /// * `ptype`: A string indicating the type of parameter ("input" or "output").
- /// * `specs`: The current list of parameter specifications (either inputs or outputs).
+ /// * `ptype`: A string indicating the type of parameter ("input" or
+ /// "output").
+ /// * `specs`: The current list of parameter specifications (either inputs
+ /// or outputs).
///
/// # Returns
///
- /// An `anyhow::Result>` containing a new vector of parameter specifications with the added parameter.
+ /// An `anyhow::Result>` containing a new vector of
+ /// parameter specifications with the added parameter.
fn with_parameter(
spec: ParameterSpec,
ptype: &str,
@@ -176,8 +195,10 @@ impl FirehoseOperatorSignature {
/// # Returns
///
/// An `Result` where:
- /// * `Ok(Self)`: A new `FirehoseOperatorSignature` with the input parameter added.
- /// * `Err(String)`: An error message if the input parameter name already exists in the signature.
+ /// * `Ok(Self)`: A new `FirehoseOperatorSignature` with the input parameter
+ /// added.
+ /// * `Err(String)`: An error message if the input parameter name already
+ /// exists in the signature.
pub fn with_input_result(
self,
spec: ParameterSpec,
@@ -222,8 +243,10 @@ impl FirehoseOperatorSignature {
/// # Returns
///
/// An `Result` where:
- /// * `Ok(Self)`: A new `FirehoseOperatorSignature` with the output parameter added.
- /// * `Err(String)`: An error message if the output parameter name already exists in the signature.
+ /// * `Ok(Self)`: A new `FirehoseOperatorSignature` with the output
+ /// parameter added.
+ /// * `Err(String)`: An error message if the output parameter name already
+ /// exists in the signature.
pub fn with_output_result(
self,
spec: ParameterSpec,
@@ -302,9 +325,12 @@ impl FirehoseOperatorSignature {
///
/// # Arguments
///
- /// * `param_type`: A string indicating the type of parameters being validated ("input" or "output").
- /// * `specs`: A slice of `ParameterSpec` that defines the expected parameters.
- /// * `provided`: A map of provided parameters, where keys are parameter names and values are their data types.
+ /// * `param_type`: A string indicating the type of parameters being
+ /// validated ("input" or "output").
+ /// * `specs`: A slice of `ParameterSpec` that defines the expected
+ /// parameters.
+ /// * `provided`: A map of provided parameters, where keys are parameter
+ /// names and values are their data types.
///
/// # Returns
///
diff --git a/crates/bimm-firehose/src/core/rows.rs b/crates/bimm-firehose/src/core/rows.rs
index 81056fdf..4b0115e7 100644
--- a/crates/bimm-firehose/src/core/rows.rs
+++ b/crates/bimm-firehose/src/core/rows.rs
@@ -1,15 +1,30 @@
-use crate::core::FirehoseValue;
-use crate::core::operations::signature::FirehoseOperatorSignature;
-use crate::core::schema::{BuildPlan, FirehoseTableSchema};
+use std::{
+ any::Any,
+ collections::HashMap,
+ fmt::Debug,
+ ops::{
+ Index,
+ IndexMut,
+ RangeBounds,
+ },
+ sync::Arc,
+ vec::Drain,
+};
+
use anyhow::Context;
-use serde::Serialize;
-use serde::de::DeserializeOwned;
-use std::any::Any;
-use std::collections::HashMap;
-use std::fmt::Debug;
-use std::ops::{Index, IndexMut, RangeBounds};
-use std::sync::Arc;
-use std::vec::Drain;
+use serde::{
+ Serialize,
+ de::DeserializeOwned,
+};
+
+use crate::core::{
+ FirehoseValue,
+ operations::signature::FirehoseOperatorSignature,
+ schema::{
+ BuildPlan,
+ FirehoseTableSchema,
+ },
+};
/// Represents a row in a Firehose table, containing values for each column.
pub struct FirehoseRow {
@@ -35,7 +50,8 @@ impl Debug for FirehoseRow {
}
impl FirehoseRow {
- /// Creates a new `ValueRow` with the given schema and initializes all slots to `None`.
+ /// Creates a new `ValueRow` with the given schema and initializes all slots
+ /// to `None`.
pub fn new(schema: Arc) -> Self {
let mut slots = Vec::with_capacity(schema.columns.len());
slots.resize_with(schema.columns.len(), || None);
@@ -93,7 +109,8 @@ pub trait FirehoseRowReader {
/// Returns the schema of the row.
fn schema(&self) -> &Arc;
- /// Returns an iterator over the column names and their corresponding values.
+ /// Returns an iterator over the column names and their corresponding
+ /// values.
fn iter(&self) -> impl Iterator- )>;
/// Returns true if the row has a value for the specified column name.
@@ -108,7 +125,8 @@ pub trait FirehoseRowReader {
///
/// # Returns
///
- /// A boolean indicating whether the row has a value for the specified column.
+ /// A boolean indicating whether the row has a value for the specified
+ /// column.
fn has_column_value(
&self,
column_name: &str,
@@ -122,8 +140,9 @@ pub trait FirehoseRowReader {
///
/// # Returns
///
- /// An `Option<&ValueBox>` containing a reference to the value of the specified column,
- /// or `None` if the column does not exist or has no value.
+ /// An `Option<&ValueBox>` containing a reference to the value of the
+ /// specified column, or `None` if the column does not exist or has no
+ /// value.
fn maybe_get(
&self,
column_name: &str,
@@ -137,7 +156,8 @@ pub trait FirehoseRowReader {
///
/// # Returns
///
- /// An `anyhow::Result<&FirehoseValue>` reference to the column value; or an error.
+ /// An `anyhow::Result<&FirehoseValue>` reference to the column value; or an
+ /// error.
fn try_get(
&self,
column_name: &str,
@@ -268,7 +288,8 @@ pub trait FirehoseRowWriter {
/// # Arguments
///
/// * `column_name`: The name of the column to set the value for.
- /// * `value`: The value to set for the specified column, wrapped in a `ValueBox`.
+ /// * `value`: The value to set for the specified column, wrapped in a
+ /// `ValueBox`.
///
/// # Panics
///
@@ -328,7 +349,8 @@ pub trait FirehoseRowWriter {
}
}
- /// Take the value of the column, setting it to `None`, and returning it as an `Option`.
+ /// Take the value of the column, setting it to `None`, and returning it as
+ /// an `Option`.
fn take_column(
&mut self,
column_name: &str,
@@ -407,7 +429,8 @@ impl Debug for FirehoseRowBatch {
}
impl FirehoseRowBatch {
- /// Creates a new `ValueRowBatch` with the given schema and initializes an empty vector of rows.
+ /// Creates a new `ValueRowBatch` with the given schema and initializes an
+ /// empty vector of rows.
pub fn new(schema: Arc) -> Self {
Self::new_with_size(schema, 0)
}
@@ -548,7 +571,8 @@ impl IndexMut for FirehoseRowBatch {
}
}
-/// A map of formal parameters to their corresponding input and output columns in a build plan.
+/// A map of formal parameters to their corresponding input and output columns
+/// in a build plan.
struct ParameterMapper {
/// The build plan that describes the operator and its inputs/outputs.
build_plan: Arc,
@@ -560,7 +584,8 @@ impl ParameterMapper {
ParameterMapper { build_plan }
}
- /// Maps an input parameter name to its corresponding column name in the build plan.
+ /// Maps an input parameter name to its corresponding column name in the
+ /// build plan.
///
/// # Arguments
///
@@ -568,7 +593,8 @@ impl ParameterMapper {
///
/// # Returns
///
- /// An `Option<&str>` containing the column name if the parameter is an input parameter,
+ /// An `Option<&str>` containing the column name if the parameter is an
+ /// input parameter,
fn try_map_input_name(
&self,
parameter_name: &str,
@@ -594,7 +620,8 @@ impl ParameterMapper {
}
}
- /// Maps an output parameter name to its corresponding column name in the build plan.
+ /// Maps an output parameter name to its corresponding column name in the
+ /// build plan.
///
/// # Arguments
///
@@ -602,7 +629,8 @@ impl ParameterMapper {
///
/// # Returns
///
- /// An `Option<&str>` containing the column name if the parameter is an output parameter,
+ /// An `Option<&str>` containing the column name if the parameter is an
+ /// output parameter,
fn try_map_output_name(
&self,
parameter_name: &str,
@@ -713,7 +741,8 @@ impl<'a> FirehoseBatchTransaction<'a> {
}
}
-/// A row-transaction for a single row in a batch, allowing both reading and writing of values.
+/// A row-transaction for a single row in a batch, allowing both reading and
+/// writing of values.
///
/// This is a view-class of a backing `FirehoseBatchTransaction`.
pub struct FirehoseRowTransaction<'a> {
@@ -779,12 +808,20 @@ impl FirehoseRowWriter for FirehoseRowTransaction<'_> {
#[cfg(test)]
mod tests {
- use super::*;
- use crate::core::operations::signature::ParameterSpec;
- use crate::core::schema::{ColumnSchema, DataTypeDescription, FirehoseTableSchema};
use std::sync::Arc;
- /// Ensures that `ValueRow` is `Send`, allowing it to be safely shared across threads.
+ use super::*;
+ use crate::core::{
+ operations::signature::ParameterSpec,
+ schema::{
+ ColumnSchema,
+ DataTypeDescription,
+ FirehoseTableSchema,
+ },
+ };
+
+ /// Ensures that `ValueRow` is `Send`, allowing it to be safely shared
+ /// across threads.
const VALUE_ROW_IS_SEND: fn() = || {
fn assert_send() {}
assert_send::();
diff --git a/crates/bimm-firehose/src/core/schema.rs b/crates/bimm-firehose/src/core/schema.rs
index 9b742ae9..bc7ab31d 100644
--- a/crates/bimm-firehose/src/core/schema.rs
+++ b/crates/bimm-firehose/src/core/schema.rs
@@ -1,8 +1,24 @@
+use std::{
+ collections::{
+ BTreeMap,
+ HashSet,
+ },
+ ops::{
+ Index,
+ IndexMut,
+ },
+};
+
+use anyhow::{
+ anyhow,
+ bail,
+};
+use serde::{
+ Deserialize,
+ Serialize,
+};
+
use crate::core::identifiers;
-use anyhow::{anyhow, bail};
-use serde::{Deserialize, Serialize};
-use std::collections::{BTreeMap, HashSet};
-use std::ops::{Index, IndexMut};
/// A serializable description of a data type.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
@@ -92,7 +108,8 @@ impl BuildPlan {
///
/// # Arguments
///
- /// - `config`: The configuration to attach to the build plan, serialized as JSON.
+ /// - `config`: The configuration to attach to the build plan, serialized as
+ /// JSON.
///
/// # Returns
///
@@ -114,7 +131,8 @@ impl BuildPlan {
///
/// # Arguments
///
- /// - `assoc`: A slice of tuples where each tuple contains a parameter name and a column name.
+ /// - `assoc`: A slice of tuples where each tuple contains a parameter name
+ /// and a column name.
///
/// # Returns
///
@@ -133,7 +151,8 @@ impl BuildPlan {
///
/// # Arguments
///
- /// - `inputs`: A slice of tuples where each tuple contains a parameter name and a column name.
+ /// - `inputs`: A slice of tuples where each tuple contains a parameter name
+ /// and a column name.
///
/// # Returns
///
@@ -152,7 +171,8 @@ impl BuildPlan {
///
/// # Arguments
///
- /// - `outputs`: A slice of tuples where each tuple contains a parameter name and a column name.
+ /// - `outputs`: A slice of tuples where each tuple contains a parameter
+ /// name and a column name.
///
/// # Returns
///
@@ -194,7 +214,8 @@ impl BuildPlan {
///
/// # Returns
///
- /// An `anyhow::Result<&str>` containing the column name corresponding to the input parameter.
+ /// An `anyhow::Result<&str>` containing the column name corresponding to
+ /// the input parameter.
pub fn translate_input_name(
&self,
parameter_name: &str,
@@ -210,7 +231,8 @@ impl BuildPlan {
///
/// # Returns
///
- /// An `anyhow::Result<&str>` containing the column name corresponding to the output parameter.
+ /// An `anyhow::Result<&str>` containing the column name corresponding to
+ /// the output parameter.
pub fn translate_output_name(
&self,
parameter_name: &str,
@@ -378,12 +400,15 @@ impl FirehoseTableSchema {
///
/// # Arguments
///
- /// - `columns`: A slice of `ColumnSchema` representing the columns in the table.
- /// - `plans`: A slice of `BuildPlan` representing the build plans for the table.
+ /// - `columns`: A slice of `ColumnSchema` representing the columns in the
+ /// table.
+ /// - `plans`: A slice of `BuildPlan` representing the build plans for the
+ /// table.
///
/// # Returns
///
- /// An `anyhow::Result<(Vec, Vec)>` containing the base columns and the ordered build plans.
+ /// An `anyhow::Result<(Vec, Vec)>` containing the base
+ /// columns and the ordered build plans.
fn check_graph(
columns: &[ColumnSchema],
plans: &[BuildPlan],
@@ -453,11 +478,13 @@ impl FirehoseTableSchema {
/// Compute the build order for the table schema.
///
- /// This function checks the build plans and their dependencies to determine the order in which they should be executed.
+ /// This function checks the build plans and their dependencies to determine
+ /// the order in which they should be executed.
///
/// # Returns
///
- /// An `anyhow::Result<(Vec, Vec)>` containing the base columns and the ordered build plans.
+ /// An `anyhow::Result<(Vec, Vec)>` containing the base
+ /// columns and the ordered build plans.
pub fn build_order(&self) -> anyhow::Result<(Vec, Vec)> {
Self::check_graph(&self.columns, &self.build_plans)
}
@@ -543,11 +570,13 @@ impl FirehoseTableSchema {
/// # Arguments
///
/// - `plan`: The build plan to add.
- /// - `output_info`: A slice of tuples where each tuple contains the output column name, its data type, and a description.
+ /// - `output_info`: A slice of tuples where each tuple contains the output
+ /// column name, its data type, and a description.
///
/// # Returns
///
- /// An `anyhow::Result<()>` indicating success or containing an error if the operation fails.
+ /// An `anyhow::Result<()>` indicating success or containing an error if the
+ /// operation fails.
pub fn add_build_plan_and_outputs(
&mut self,
plan: BuildPlan,
@@ -576,7 +605,8 @@ impl FirehoseTableSchema {
///
/// # Returns
///
- /// An `anyhow::Result<()>` indicating success or containing an error if the operation fails.
+ /// An `anyhow::Result<()>` indicating success or containing an error if the
+ /// operation fails.
pub fn extend_via_plan(
&mut self,
plan: BuildPlan,
@@ -652,7 +682,8 @@ impl FirehoseTableSchema {
///
/// # Returns
///
- /// An `Option` containing the index of the column if found, or `None` if not found.
+ /// An `Option` containing the index of the column if found, or
+ /// `None` if not found.
pub fn column_index(
&self,
name: &str,
@@ -673,7 +704,8 @@ impl FirehoseTableSchema {
///
/// # Returns
///
- /// An `anyhow::Result` containing the index of the column if it exists.
+ /// An `anyhow::Result` containing the index of the column if it
+ /// exists.
pub fn check_column_index(
&self,
name: &str,
@@ -687,9 +719,10 @@ impl FirehoseTableSchema {
#[cfg(test)]
mod tests {
- use super::*;
use indoc::indoc;
+ use super::*;
+
/// Ensures that `FirehoseTableSchema` is `Send`.
const FIREHOSE_TABLE_SCHEMA_IS_SEND: fn() = || {
fn assert_send() {}
diff --git a/crates/bimm-firehose/src/core/values.rs b/crates/bimm-firehose/src/core/values.rs
index 3c049b76..61b3e88c 100644
--- a/crates/bimm-firehose/src/core/values.rs
+++ b/crates/bimm-firehose/src/core/values.rs
@@ -1,10 +1,22 @@
-use anyhow::{Context, bail};
-use serde::Serialize;
-use serde::de::DeserializeOwned;
-use std::any::{Any, TypeId};
-use std::fmt::Debug;
-
-/// A wrapper type that can hold either a JSON value or a boxed value of any type.
+use std::{
+ any::{
+ Any,
+ TypeId,
+ },
+ fmt::Debug,
+};
+
+use anyhow::{
+ Context,
+ bail,
+};
+use serde::{
+ Serialize,
+ de::DeserializeOwned,
+};
+
+/// A wrapper type that can hold either a JSON value or a boxed value of any
+/// type.
pub enum FirehoseValue {
/// Holds a JSON value.
Value(serde_json::Value),
@@ -87,7 +99,8 @@ impl FirehoseValue {
FirehoseValue::Value(value)
}
- /// Creates a new `ValueBox::Boxed` by boxing an object that implements `Any` and `Send`.
+ /// Creates a new `ValueBox::Boxed` by boxing an object that implements
+ /// `Any` and `Send`.
pub fn boxing(obj: T) -> Self
where
T: Any + 'static + Send,
@@ -95,7 +108,8 @@ impl FirehoseValue {
Self::from_box(Box::new(obj))
}
- /// Creates a new `ValueBox::Boxed` from a boxed object that implements `Any` and `Send`.
+ /// Creates a new `ValueBox::Boxed` from a boxed object that implements
+ /// `Any` and `Send`.
pub fn from_box(boxed: Box) -> Self
where
T: Any + 'static + Send,
@@ -114,7 +128,8 @@ impl FirehoseValue {
matches!(self, FirehoseValue::Boxed(_))
}
- /// Unwraps the `ValueBox` and returns a reference to the contained JSON value.
+ /// Unwraps the `ValueBox` and returns a reference to the contained JSON
+ /// value.
pub fn unwrap_value(&self) -> &serde_json::Value {
if let FirehoseValue::Value(value) = self {
value
@@ -127,7 +142,8 @@ impl FirehoseValue {
///
/// # Type Parameters
///
- /// `T`: The type to deserialize the JSON value into. It must implement `DeserializeOwned`.
+ /// `T`: The type to deserialize the JSON value into. It must implement
+ /// `DeserializeOwned`.
///
/// # Panics
///
@@ -154,7 +170,8 @@ impl FirehoseValue {
///
/// # Type Parameters
///
- /// `T`: The type to deserialize the JSON value into. It must implement `DeserializeOwned`.
+ /// `T`: The type to deserialize the JSON value into. It must implement
+ /// `DeserializeOwned`.
///
/// # Panics
///
@@ -171,7 +188,8 @@ impl FirehoseValue {
self.parse_as::().unwrap()
}
- /// Unwraps the `ValueBox` and returns a reference to the contained boxed value.
+ /// Unwraps the `ValueBox` and returns a reference to the contained boxed
+ /// value.
///
/// # Type Parameters
///
@@ -201,7 +219,8 @@ impl FirehoseValue {
}
}
- /// Unwraps the `ValueBox` and returns a reference to the contained boxed value.
+ /// Unwraps the `ValueBox` and returns a reference to the contained boxed
+ /// value.
///
/// # Type Parameters
///
@@ -209,7 +228,8 @@ impl FirehoseValue {
///
/// # Panics
///
- /// If the `ValueBox` does not contain a boxed value or if the downcast fails.
+ /// If the `ValueBox` does not contain a boxed value or if the downcast
+ /// fails.
///
/// # Returns
///
diff --git a/crates/bimm-firehose/src/lib.rs b/crates/bimm-firehose/src/lib.rs
index 4942b8ed..69d50ae8 100644
--- a/crates/bimm-firehose/src/lib.rs
+++ b/crates/bimm-firehose/src/lib.rs
@@ -18,12 +18,13 @@ mod pipeline;
/// Define a self-referential ID.
///
-/// The id will be defined as a static string constant that refers to its own namespace path.
+/// The id will be defined as a static string constant that refers to its own
+/// namespace path.
///
/// # Arguments
///
-/// * `$name`: The final path name of the ID to define;
-/// the rest of the name will be taken from the module context.
+/// * `$name`: The final path name of the ID to define; the rest of the name
+/// will be taken from the module context.
///
/// # Example
/// ```
@@ -31,7 +32,6 @@ mod pipeline;
/// bimm_firehose::define_self_referential_id!(ID);
/// // pub static ID: &str = "foo::bar::ID";
/// ```
-///
#[macro_export]
macro_rules! define_self_referential_id {
($name:ident) => {
diff --git a/crates/bimm-firehose/src/ops/mod.rs b/crates/bimm-firehose/src/ops/mod.rs
index b3f5013e..f1f1eb08 100644
--- a/crates/bimm-firehose/src/ops/mod.rs
+++ b/crates/bimm-firehose/src/ops/mod.rs
@@ -1,5 +1,7 @@
-use crate::core::operations::environment::MapOpEnvironment;
-use crate::core::operations::registration::FirehoseOperatorFactoryRegistration;
+use crate::core::operations::{
+ environment::MapOpEnvironment,
+ registration::FirehoseOperatorFactoryRegistration,
+};
/// Build the default environment.
///
diff --git a/crates/bimm-firehose/src/pipeline/data_load_operator.rs b/crates/bimm-firehose/src/pipeline/data_load_operator.rs
index 38bc4ce5..4fd9ad9a 100644
--- a/crates/bimm-firehose/src/pipeline/data_load_operator.rs
+++ b/crates/bimm-firehose/src/pipeline/data_load_operator.rs
@@ -1,6 +1,7 @@
-use crate::pipeline::DataLoadMetaDataItem;
use std::fmt::Debug;
+use crate::pipeline::DataLoadMetaDataItem;
+
/// Support super-trait for loadable data types.
///
/// # Trait Requirements
@@ -10,7 +11,8 @@ use std::fmt::Debug;
/// - `Send` and `Sync`: the data item must be thread-safe.
pub trait DataLoadDataItem: Debug + Clone + Send + Sync {}
-/// Blanket implementation of `DataLoadDataItem` for any type that meets the requirements.
+/// Blanket implementation of `DataLoadDataItem` for any type that meets the
+/// requirements.
impl DataLoadDataItem for T where T: Debug + Clone + Send + Sync {}
/// Trait for data load operators.
@@ -23,7 +25,8 @@ where
///
/// # Arguments
///
- /// * `meta` - Metadata item that contains information needed to load the data.
+ /// * `meta` - Metadata item that contains information needed to load the
+ /// data.
///
/// # Returns
///
@@ -46,7 +49,8 @@ where
/// The function that defines how to load the data.
func: F,
- /// Phantom data to associate the operator with specific metadata and data types.
+ /// Phantom data to associate the operator with specific metadata and data
+ /// types.
phantom: std::marker::PhantomData<(M, T)>,
}
diff --git a/crates/bimm-firehose/src/pipeline/data_load_plan.rs b/crates/bimm-firehose/src/pipeline/data_load_plan.rs
index f04ff2e1..d234e39f 100644
--- a/crates/bimm-firehose/src/pipeline/data_load_plan.rs
+++ b/crates/bimm-firehose/src/pipeline/data_load_plan.rs
@@ -1,7 +1,14 @@
-use crate::pipeline::{DataLoadDataItem, DataLoadMetaDataItem, DataLoadOperator, DataLoadSchedule};
use std::sync::Arc;
-/// Represents a plan for loading data, including a schedule and an optional operator.
+use crate::pipeline::{
+ DataLoadDataItem,
+ DataLoadMetaDataItem,
+ DataLoadOperator,
+ DataLoadSchedule,
+};
+
+/// Represents a plan for loading data, including a schedule and an optional
+/// operator.
#[derive(Debug, Clone)]
pub struct DataLoadPlan
where
@@ -28,7 +35,8 @@ where
}
}
- /// Initializes a `DataLoadPlan` with a given schedule and an optional operator.
+ /// Initializes a `DataLoadPlan` with a given schedule and an optional
+ /// operator.
pub fn init(
schedule: DataLoadSchedule,
op: Option>>,
@@ -49,9 +57,10 @@ where
#[cfg(test)]
mod tests {
+ use std::sync::Arc;
+
use super::*;
use crate::pipeline::FnOperator;
- use std::sync::Arc;
#[test]
fn test_index_to_str_plan() {
diff --git a/crates/bimm-firehose/src/pipeline/data_load_schedule.rs b/crates/bimm-firehose/src/pipeline/data_load_schedule.rs
index 7136fb70..7d059c7e 100644
--- a/crates/bimm-firehose/src/pipeline/data_load_schedule.rs
+++ b/crates/bimm-firehose/src/pipeline/data_load_schedule.rs
@@ -1,21 +1,29 @@
-use serde::{Deserialize, Serialize};
-use std::fmt::Debug;
-use std::hash::Hash;
+use std::{
+ fmt::Debug,
+ hash::Hash,
+};
+
+use serde::{
+ Deserialize,
+ Serialize,
+};
/// A fixed schedule for loading data items.
///
-/// The type `M` must implement the `ScheduleItem` trait, which requires it to be
-/// Debug, Serialize, Sync, and Send safe.
+/// The type `M` must implement the `ScheduleItem` trait, which requires it to
+/// be Debug, Serialize, Sync, and Send safe.
///
/// # Motivation
///
-/// This structure is designed to hold the intermediate schedule for a data pipeline
-/// load operation; in a format which can be serialized and deserialized;
-/// and can be transformed into modified schedules without forcing a data load.
+/// This structure is designed to hold the intermediate schedule for a data
+/// pipeline load operation; in a format which can be serialized and
+/// deserialized; and can be transformed into modified schedules without forcing
+/// a data load.
///
/// As such, it is a thin wrapper around a `Vec`, where `M` carries complex
/// constraints; but the expectation is that additional metadata will be added
-/// to the schedule in the future (such as block size, or item load cost estimates).
+/// to the schedule in the future (such as block size, or item load cost
+/// estimates).
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(bound = "M: DataLoadMetaDataItem")]
pub struct DataLoadSchedule
@@ -28,9 +36,9 @@ where
/// Support super-trait for `DataLoadSchedule` items.
///
-/// This trait exists to simplify the requirements for items in a `DataLoadSchedule`;
-/// blanket implementations of `ScheduleItem` are provided for any type that meets
-/// the requirements.
+/// This trait exists to simplify the requirements for items in a
+/// `DataLoadSchedule`; blanket implementations of `ScheduleItem` are provided
+/// for any type that meets the requirements.
///
/// # Trait Requirements
///
@@ -38,13 +46,15 @@ where
/// - `Clone`: the schedule must be cloneable.
/// - `PartialEq` and `Eq`: the schedule must be comparable for equality.
/// - `Send` and `Sync`: the schedule must be thread-safe.
-/// - `Serialize` and `Deserialize`: the schedule must be serializable and deserializable.
+/// - `Serialize` and `Deserialize`: the schedule must be serializable and
+/// deserializable.
pub trait DataLoadMetaDataItem:
Debug + Clone + Hash + PartialEq + Eq + Send + Sync + Serialize + for<'de> Deserialize<'de>
{
}
-/// Blanket implementation of `ScheduleItem` for any type that meets the requirements.
+/// Blanket implementation of `ScheduleItem` for any type that meets the
+/// requirements.
impl DataLoadMetaDataItem for T where
T: Debug + Clone + Hash + PartialEq + Eq + Send + Sync + Serialize + for<'de> Deserialize<'de>
{
@@ -95,11 +105,13 @@ where
///
/// # Arguments
///
- /// - `predicate`: A closure that takes a reference to an item of type `M` and returns a boolean.
+ /// - `predicate`: A closure that takes a reference to an item of type `M`
+ /// and returns a boolean.
///
/// # Returns
///
- /// A new `DataLoadSchedule` containing only the items that match the predicate.
+ /// A new `DataLoadSchedule` containing only the items that match the
+ /// predicate.
pub fn filter
(
&self,
predicate: P,
@@ -121,15 +133,18 @@ where
/// Filters the items in the schedule based on a provided filter function,
/// and maps them to a new type `R` if the predicate returns `Some(R)`.
///
- /// This function is useful for transforming the schedule items while filtering them.
+ /// This function is useful for transforming the schedule items while
+ /// filtering them.
///
/// # Arguments
///
- /// - `predicate`: A closure that takes a reference to an item of type `M` and returns an `Option`.
+ /// - `predicate`: A closure that takes a reference to an item of type `M`
+ /// and returns an `Option`.
///
/// # Returns
///
- /// A new `DataLoadSchedule` containing the items that matched the predicate and were transformed to type `R`.
+ /// A new `DataLoadSchedule` containing the items that matched the
+ /// predicate and were transformed to type `R`.
pub fn filter_map(
&self,
predicate: P,
diff --git a/crates/bimm-firehose/src/pipeline/data_schedule_source.rs b/crates/bimm-firehose/src/pipeline/data_schedule_source.rs
index 275af872..e0c59c23 100644
--- a/crates/bimm-firehose/src/pipeline/data_schedule_source.rs
+++ b/crates/bimm-firehose/src/pipeline/data_schedule_source.rs
@@ -1,6 +1,10 @@
-use crate::pipeline::{DataLoadMetaDataItem, DataLoadSchedule};
use std::fmt::Debug;
+use crate::pipeline::{
+ DataLoadMetaDataItem,
+ DataLoadSchedule,
+};
+
/// Data-pipeline trait for schedule sources.
pub trait DataScheduleSource: Debug
where
@@ -21,7 +25,8 @@ where
}
}
-/// Blanket implementation of `DataScheduleSourceExt` for any type that meets the requirements.
+/// Blanket implementation of `DataScheduleSourceExt` for any type that meets
+/// the requirements.
impl DataScheduleSourceExt for S
where
M: DataLoadMetaDataItem,
@@ -66,7 +71,8 @@ where
}
}
-/// Demo of a simple filter source that filters items based on a predicate function.
+/// Demo of a simple filter source that filters items based on a predicate
+/// function.
pub struct SimpleFilterSource
where
M: DataLoadMetaDataItem,
@@ -96,7 +102,8 @@ impl SimpleFilterSource
where
M: DataLoadMetaDataItem,
{
- /// Creates a new `SimpleFilterSource` with the provided inner source and predicate.
+ /// Creates a new `SimpleFilterSource` with the provided inner source and
+ /// predicate.
///
/// # Arguments
///
@@ -121,7 +128,8 @@ where
}
}
-/// A wrapper that maps the output of a schedule source to a different type using a mapping function.
+/// A wrapper that maps the output of a schedule source to a different type
+/// using a mapping function.
pub struct ScheduleSourceMappingWrapper
where
A: DataLoadMetaDataItem,
@@ -157,7 +165,8 @@ where
B: DataLoadMetaDataItem,
F: Fn(&DataLoadSchedule) -> anyhow::Result> + Send + Sync + 'static,
{
- /// Creates a new `ScheduleSourceMappingWrapper` with the provided inner source and mapping function.
+ /// Creates a new `ScheduleSourceMappingWrapper` with the provided inner
+ /// source and mapping function.
///
/// # Arguments
///
diff --git a/crates/bimm-firehose/src/utility/probability.rs b/crates/bimm-firehose/src/utility/probability.rs
index 76ebb30a..1d04a73e 100644
--- a/crates/bimm-firehose/src/utility/probability.rs
+++ b/crates/bimm-firehose/src/utility/probability.rs
@@ -1,6 +1,7 @@
+use std::fmt::Debug;
+
use anyhow::bail;
use num_traits::Float;
-use std::fmt::Debug;
/// Validate a probability in the range ``[0.0, 1.0]``.
///
diff --git a/crates/bimm/Cargo.toml b/crates/bimm/Cargo.toml
index 4b096664..04c337dc 100644
--- a/crates/bimm/Cargo.toml
+++ b/crates/bimm/Cargo.toml
@@ -12,12 +12,11 @@ description = "burn image models"
workspace = true
[dependencies]
-bimm-contracts = { workspace = true, features = ["burn"] }
+bunsen = { workspace = true, features = ["cache"] }
+burn = { workspace = true, features = ["network", "store"] }
+burn-store = { workspace = true, features = ["pytorch"] }
-# There is a dep/gate bug that requires "ndarray" when using "network".
-burn = { workspace = true, features = ["network", "ndarray"] }
dirs = { workspace = true }
-burn-import = { workspace = true }
crc = { workspace = true }
serde = { workspace = true, features = ["derive"] }
@@ -27,7 +26,7 @@ anyhow = { workspace = true }
num-traits = { workspace = true }
[dev-dependencies]
-burn = { workspace = true, features = ["autodiff", "wgpu"] }
+burn = { workspace = true, features = ["autodiff"] }
hamcrest = { workspace = true }
indoc = { workspace = true }
criterion = { workspace = true }
@@ -36,6 +35,7 @@ criterion = { workspace = true }
[features]
default = []
std = []
+cuda = ["burn/cuda", "bunsen/cuda"]
[[bench]]
name = "drop_block"
diff --git a/crates/bimm/benches/drop_block.rs b/crates/bimm/benches/drop_block.rs
index cd2b83b1..653dbdd0 100644
--- a/crates/bimm/benches/drop_block.rs
+++ b/crates/bimm/benches/drop_block.rs
@@ -1,12 +1,22 @@
-use bimm::layers::drop::drop_block::{DropBlockOptions, drop_block_2d};
-use bimm::utility::burn::noise::NoiseConfig;
-use burn::backend::NdArray;
-use burn::prelude::Tensor;
-use criterion::{Criterion, criterion_group, criterion_main};
use std::hint::black_box;
+use bunsen::{
+ blocks::images::drop::drop_block::{
+ DropBlockOptions,
+ drop_block_2d,
+ },
+ ops::noise::NoiseConfig,
+ support::testing::PerfTestBackend,
+};
+use burn::prelude::Tensor;
+use criterion::{
+ Criterion,
+ criterion_group,
+ criterion_main,
+};
+
fn bench_drop_block_10x32x32x3_7_normalise(c: &mut Criterion) {
- type B = NdArray;
+ type B = PerfTestBackend;
let device = Default::default();
let batch_size = 10;
diff --git a/crates/bimm/src/cache/disk.rs b/crates/bimm/src/cache/disk.rs
deleted file mode 100644
index 5202901d..00000000
--- a/crates/bimm/src/cache/disk.rs
+++ /dev/null
@@ -1,114 +0,0 @@
-//! # Cache Policy
-
-use alloc::string::String;
-use anyhow::bail;
-use burn::config::Config;
-use burn::data::network::downloader;
-use std::fs::{File, remove_file};
-use std::io::Write;
-use std::path::PathBuf;
-
-/// Cache Policy
-#[derive(Config, Debug)]
-pub struct DiskCacheConfig {
- /// Key for the root cache directory.
- #[config(default = "\"bimm\".to_string()")]
- pub root_cache_key: String,
-}
-
-impl Default for DiskCacheConfig {
- fn default() -> Self {
- Self::new()
- }
-}
-
-impl DiskCacheConfig {
- /// Fetch the base cache directory.
- ///
- /// If the cache directory does not exist, does not create it.
- pub fn base_cache_dir(&self) -> anyhow::Result {
- Ok(dirs::home_dir()
- .expect("Should be able to get home directory")
- .join(".cache")
- .join(&self.root_cache_key))
- }
-
- /// Fetch the base cache directory.
- ///
- /// If the cache directory does not exist, creates it.
- pub fn ensure_base_cache_dir(&self) -> anyhow::Result {
- let dir = self.base_cache_dir()?;
- if !dir.exists() {
- std::fs::create_dir_all(&dir)?;
- }
- Ok(dir)
- }
-
- /// Map a resource key to a cache path.
- ///
- /// Does not ensure that the path (or any of the parents) exist.
- pub fn resource_to_path(
- &self,
- resource_key: &[String],
- ) -> anyhow::Result {
- let path = self.base_cache_dir()?;
- Ok(resource_key.iter().fold(path, |acc, s| acc.join(s)))
- }
-
- /// Map a resource key to a cache path and ensure the parent directory exists.
- pub fn ensure_resource_parent_dir(
- &self,
- resource_key: &[String],
- ) -> anyhow::Result {
- let path = self.resource_to_path(resource_key)?;
- if !path.exists() {
- std::fs::create_dir_all(path.parent().unwrap())?;
- }
- Ok(path)
- }
-
- /// Fetch a Resource to the Cache.
- pub fn fetch_resource(
- &self,
- url: &str,
- resource: &[String],
- ) -> anyhow::Result {
- let cache_file_path = self.ensure_resource_parent_dir(resource)?;
- try_cache_download_to_path(url, cache_file_path)
- }
-}
-
-/// Download a URL resource to a given path.
-///
-/// If the path already exists, does nothing.
-///
-/// # Returns
-///
-/// The cache path.
-pub fn try_cache_download_to_path(
- url: &str,
- cache_file_path: PathBuf,
-) -> anyhow::Result {
- if !cache_file_path.exists() {
- let file_name = cache_file_path
- .file_name()
- .unwrap()
- .to_string_lossy()
- .to_string();
-
- // TODO: download-to-file instead of download-to-memory.
- // Download file content
- let bytes = downloader::download_file_as_bytes(url, &file_name);
-
- // Write content to file
- let mut output_file = File::create(&cache_file_path)?;
- let bytes_written = output_file.write(&bytes)?;
-
- if bytes_written != bytes.len() {
- remove_file(cache_file_path)?;
- bail!("Failed to write the whole model weights file.");
- }
- }
-
- Ok(cache_file_path)
-}
diff --git a/crates/bimm/src/cache/mod.rs b/crates/bimm/src/cache/mod.rs
index e01ce5a8..e7b798ca 100644
--- a/crates/bimm/src/cache/mod.rs
+++ b/crates/bimm/src/cache/mod.rs
@@ -1,6 +1,18 @@
//! # Cache Implementations
-pub mod disk;
-pub mod weights;
+use bunsen::cache as bcache;
-pub mod prefabs;
+mod prefabs;
+pub use bcache::{
+ DiskCacheConfig,
+ PretrainedWeightsDescriptor,
+ PretrainedWeightsMap,
+ StaticPretrainedWeightsDescriptor,
+ StaticPretrainedWeightsMap,
+ fetch_model_weights,
+ pretrained_weights_resource_key,
+ try_cache_download_to_path,
+ url_to_cache_key,
+};
+#[doc(inline)]
+pub use prefabs::*;
diff --git a/crates/bimm/src/cache/prefabs.rs b/crates/bimm/src/cache/prefabs.rs
index 2427fc80..496f6639 100644
--- a/crates/bimm/src/cache/prefabs.rs
+++ b/crates/bimm/src/cache/prefabs.rs
@@ -1,15 +1,24 @@
//! # Config Prefabs for Well-Known Model Configurations
-use crate::cache::weights::{
- PretrainedWeightsDescriptor, PretrainedWeightsMap, StaticPretrainedWeightsMap,
+use alloc::{
+ collections::BTreeMap,
+ format,
+ string::{
+ String,
+ ToString,
+ },
+ sync::Arc,
};
-use alloc::collections::BTreeMap;
-use alloc::format;
-use alloc::string::{String, ToString};
-use alloc::sync::Arc;
+use core::fmt::Debug;
+
use anyhow::bail;
use burn::config::Config;
-use core::fmt::Debug;
+
+use crate::cache::{
+ PretrainedWeightsDescriptor,
+ PretrainedWeightsMap,
+ StaticPretrainedWeightsMap,
+};
/// Static builder for a [`PreFabConfig`]
pub struct StaticPreFabConfig
diff --git a/crates/bimm/src/cache/weights.rs b/crates/bimm/src/cache/weights.rs
deleted file mode 100644
index 749a585e..00000000
--- a/crates/bimm/src/cache/weights.rs
+++ /dev/null
@@ -1,227 +0,0 @@
-//! # Module / Weight Caches
-
-use crate::cache::disk::DiskCacheConfig;
-use alloc::collections::BTreeMap;
-use alloc::string::{String, ToString};
-use alloc::vec::Vec;
-use alloc::{format, vec};
-use anyhow::bail;
-use serde::{Deserialize, Serialize};
-use std::path::PathBuf;
-
-const X25: crc::Crc = crc::Crc::::new(&crc::CRC_16_IBM_SDLC);
-
-/// Returns a local path to model weights file.
-/// If the file does not exist, it will be downloaded from the given URL.
-pub fn fetch_model_weights>(url: S) -> anyhow::Result {
- let cache_key = url_to_cache_key(Some("model"), url.as_ref());
- let resource = pretrained_weights_resource_key(&cache_key);
-
- let disk_cache = DiskCacheConfig::default();
- disk_cache.fetch_resource(url.as_ref(), &resource)
-}
-
-/// Build a cache key (bare cache file name) from a name and URL.
-pub fn url_to_cache_key(
- name: Option<&str>,
- url: &str,
-) -> String {
- let hash = X25.checksum(url.as_bytes()).to_string();
- let base_name = url.rsplit_once('/').unwrap().1;
- match name {
- Some(n) => format!("{}-{}-{}", n, hash, base_name),
- None => format!("{}-{}", hash, base_name),
- }
-}
-
-/// Get the cache resource key for a pretrained weights file.
-///
-/// # Arguments
-///
-/// - `cache_key`: the cache key (the bare cache file name).
-///
-/// # Returns
-///
-/// The cache resource key.
-pub fn pretrained_weights_resource_key(cache_key: &str) -> Vec {
- vec!["weights".to_string(), cache_key.to_string()]
-}
-
-/// Static [`PretrainedWeightsDescriptor`] provider.
-#[derive(Debug)]
-pub struct StaticPretrainedWeightsDescriptor<'a> {
- /// Name of the model.
- pub name: &'a str,
-
- /// Description of the model.
- pub description: &'a str,
-
- /// License.
- pub license: Option<&'a str>,
-
- /// Source URL.
- pub origin: Option<&'a str>,
-
- /// URL to download the weights from.
- pub urls: &'a [&'a str],
-}
-
-impl<'a> StaticPretrainedWeightsDescriptor<'a> {
- /// Convert to a [`PretrainedWeightsDescriptor`].
- pub fn to_descriptor(&self) -> PretrainedWeightsDescriptor {
- PretrainedWeightsDescriptor {
- name: self.name.to_string(),
- description: self.description.to_string(),
- license: self.license.map(|s| s.to_string()),
- origin: self.origin.map(|s| s.to_string()),
- urls: self.urls.iter().map(|s| s.to_string()).collect(),
- }
- }
-}
-
-impl From<&StaticPretrainedWeightsDescriptor<'_>> for PretrainedWeightsDescriptor {
- fn from(descriptor: &StaticPretrainedWeightsDescriptor) -> Self {
- descriptor.to_descriptor()
- }
-}
-
-/// A descriptor for a pretrained weights file.
-#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
-pub struct PretrainedWeightsDescriptor {
- /// Name of the model.
- pub name: String,
-
- /// Description of the model.
- pub description: String,
-
- /// License.
- pub license: Option,
-
- /// Source URL.
- pub origin: Option,
-
- /// URL to download the weights from.
- pub urls: Vec,
-}
-
-impl PretrainedWeightsDescriptor {
- /// Cache Key
- ///
- /// The key is ``{name}-{url crc hash}-{url basename}``.
- pub fn cache_key(&self) -> String {
- url_to_cache_key(Some(&self.name), self.urls.first().unwrap())
- }
-
- /// Read-Through Cache the Model Weights
- ///
- /// # Returns
- ///
- /// The disk location of the cached weights.
- pub fn fetch_weights(
- &self,
- disk_cache: &DiskCacheConfig,
- ) -> anyhow::Result {
- let url = self.urls.first().unwrap();
- let cache_key = &self.cache_key();
- let resource = pretrained_weights_resource_key(cache_key);
-
- disk_cache.fetch_resource(url, &resource)
- }
-}
-
-/// Static [`PretrainedWeightsMap`] builder.
-#[derive(Debug)]
-pub struct StaticPretrainedWeightsMap<'a> {
- /// List of static descriptors.
- pub items: &'a [&'a StaticPretrainedWeightsDescriptor<'a>],
-}
-
-impl<'a> StaticPretrainedWeightsMap<'a> {
- /// Convert to a [`PretrainedWeightsMap`].
- pub fn to_directory(&self) -> PretrainedWeightsMap {
- PretrainedWeightsMap {
- items: self
- .items
- .iter()
- .map(|d| {
- let desc = d.to_descriptor();
- (desc.name.clone(), desc)
- })
- .collect(),
- }
- }
-}
-
-impl<'a> From<&StaticPretrainedWeightsMap<'a>> for PretrainedWeightsMap {
- fn from(directory: &StaticPretrainedWeightsMap) -> Self {
- directory.to_directory()
- }
-}
-
-/// Directory of [`PretrainedWeightsDescriptor`]s.
-#[derive(Debug, Clone)]
-pub struct PretrainedWeightsMap {
- /// Map of descriptors.
- pub items: BTreeMap,
-}
-
-impl PretrainedWeightsMap {
- /// Lookup a descriptor by name.
- pub fn lookup_by_name(
- &self,
- name: &str,
- ) -> Option {
- self.items.get(name).cloned()
- }
-
- /// Lookup a descriptor.
- pub fn try_lookup_by_name(
- &self,
- name: &str,
- ) -> anyhow::Result {
- match self.lookup_by_name(name) {
- Some(d) => Ok(d),
- None => bail!("Descriptor not found: {}", name),
- }
- }
-
- /// Lookup a descriptor.
- pub fn expect_lookup_by_name(
- &self,
- name: &str,
- ) -> PretrainedWeightsDescriptor {
- match self.try_lookup_by_name(name) {
- Ok(p) => p,
- Err(e) => panic!("{}", e),
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use alloc::string::{String, ToString};
-
- #[test]
- fn test_static_descriptor_to_descriptor() {
- let s_desc = StaticPretrainedWeightsDescriptor {
- name: "my_model",
- description: "some description of my model.",
- urls: &["foo", "bar"],
- license: Some("MIT"),
- origin: Some("https://github.com/my_org/my_model"),
- };
- let d_desc = s_desc.to_descriptor();
-
- assert_eq!(d_desc.name, s_desc.name.to_string());
- assert_eq!(d_desc.description, s_desc.description.to_string());
- assert_eq!(
- d_desc.urls,
- s_desc
- .urls
- .iter()
- .map(|s| s.to_string())
- .collect::>()
- );
- }
-}
diff --git a/crates/bimm/src/compat/conv_shape.rs b/crates/bimm/src/compat/conv_shape.rs
index 40a5c8e6..cfa0487f 100644
--- a/crates/bimm/src/compat/conv_shape.rs
+++ b/crates/bimm/src/compat/conv_shape.rs
@@ -21,7 +21,8 @@ use alloc::vec::Vec;
/// - `input_size`: The input dimension size, must be > 0.
/// - `kernel_size`: The kernel size, must be > 0.
/// - `stride`: The stride of the convolution, must be > 0.
-/// - `padding`: The padding of the convolution, added evenly to all sides of the input.
+/// - `padding`: The padding of the convolution, added evenly to all sides of
+/// the input.
/// - `dilation`: The dilation of the convolution, must be > 0.
///
/// # Returns
@@ -69,7 +70,8 @@ pub fn maybe_conv1d_output_size(
/// - `input_size`: The input dimension size, must be > 0.
/// - `kernel_size`: The kernel size, must be > 0.
/// - `stride`: The stride of the convolution, must be > 0.
-/// - `padding`: The padding of the convolution, added evenly to all sides of the input.
+/// - `padding`: The padding of the convolution, added evenly to all sides of
+/// the input.
/// - `dilation`: The dilation of the convolution, must be > 0.
///
/// # Returns
@@ -101,10 +103,14 @@ pub fn expect_conv1d_output_size(
/// # Arguments
///
/// - `input_shape`: The input dimension shape, each dim must be > 0.
-/// - `kernel_shape`: The kernel shape; length must match `input_shape`, each dim must be > 0.
-/// - `stride`: The stride of the convolution; length must match `input_shape`, each dim must be > 0.
-/// - `padding`: The padding of the convolution, added evenly to all sides of the input.
-/// - `dilation`: The dilation of the convolution; length must match `input_shape`, each dim must be > 0.
+/// - `kernel_shape`: The kernel shape; length must match `input_shape`, each
+/// dim must be > 0.
+/// - `stride`: The stride of the convolution; length must match `input_shape`,
+/// each dim must be > 0.
+/// - `padding`: The padding of the convolution, added evenly to all sides of
+/// the input.
+/// - `dilation`: The dilation of the convolution; length must match
+/// `input_shape`, each dim must be > 0.
///
/// # Returns
///
@@ -145,7 +151,8 @@ pub fn maybe_conv_output_shape_dyn(
/// - `input_shape`: The input dimension shape, each dim must be > 0.
/// - `kernel_shape`: The kernel shape, each dim must be > 0.
/// - `stride`: The stride of the convolution, each dim must be > 0.
-/// - `padding`: The padding of the convolution, added evenly to all sides of the input.
+/// - `padding`: The padding of the convolution, added evenly to all sides of
+/// the input.
/// - `dilation`: The dilation of the convolution, each dim must be > 0.
///
/// # Returns
@@ -175,7 +182,8 @@ pub fn expect_conv_output_shape_dyn(
/// - `input_shape`: The input dimension shape, each dim must be > 0.
/// - `kernel_shape`: The kernel shape, each dim must be > 0.
/// - `stride`: The stride of the convolution, each dim must be > 0.
-/// - `padding`: The padding of the convolution, added evenly to all sides of the input.
+/// - `padding`: The padding of the convolution, added evenly to all sides of
+/// the input.
/// - `dilation`: The dilation of the convolution, each dim must be > 0.
///
/// # Returns
@@ -211,7 +219,8 @@ pub fn maybe_conv_output_shape(
/// - `input_shape`: The input dimension shape, each dim must be > 0.
/// - `kernel_shape`: The kernel shape, each dim must be > 0.
/// - `stride`: The stride of the convolution, each dim must be > 0.
-/// - `padding`: The padding of the convolution, added evenly to all sides of the input.
+/// - `padding`: The padding of the convolution, added evenly to all sides of
+/// the input.
/// - `dilation`: The dilation of the convolution, each dim must be > 0.
///
/// # Returns
diff --git a/crates/bimm/src/compat/mod.rs b/crates/bimm/src/compat/mod.rs
index 23625135..87f0c23a 100644
--- a/crates/bimm/src/compat/mod.rs
+++ b/crates/bimm/src/compat/mod.rs
@@ -4,4 +4,4 @@
//! module to link against multiple versions of `burn_support`.
pub mod conv_shape;
-pub mod ops;
+pub mod type_mapper;
diff --git a/crates/bimm/src/compat/ops.rs b/crates/bimm/src/compat/ops.rs
deleted file mode 100644
index 1102cb00..00000000
--- a/crates/bimm/src/compat/ops.rs
+++ /dev/null
@@ -1,213 +0,0 @@
-//! [`Tensor`] ops.
-
-use alloc::vec;
-use alloc::vec::Vec;
-use burn::prelude::{Backend, Tensor};
-use std::f64;
-
-/// Create a vector with evenly spaced floating point values.
-///
-/// This function generates a vector starting from `start`, ending at `end`, and incrementing by `step`.
-///
-/// # Parameters
-///
-/// - `start`: The starting value of the range.
-/// - `end`: The end value of the range (exclusive).
-/// - `step`: An optional step value. If not provided, defaults to `1.0` if `start < end`, or `-1.0` if `start > end`.
-///
-/// # Returns
-///
-/// A vector containing the generated floating point values.
-#[must_use]
-pub fn float_vec_arange(
- start: f64,
- end: f64,
- step: Option,
-) -> Vec {
- assert_ne!(start, end);
- let step = if start < end {
- let step = step.unwrap_or(1.0);
- if step <= 0.0 {
- panic!("Step must be positive when start < end");
- }
- step
- } else {
- let step = step.unwrap_or(-1.0);
- if step >= 0.0 {
- panic!("Step must be negative when start > end");
- }
- step
- };
-
- let mut values: Vec = Vec::new();
- loop {
- let acc = start + values.len() as f64 * step;
- if (step > 0.0 && acc > end) || (step < 0.0 && acc < end) {
- break;
- }
- values.push(acc);
- }
-
- values
-}
-
-/// Create a vector with evenly spaced floating point values.
-///
-/// This function generates a vector with `num` values starting from `start`, ending at `end`, and evenly spaced.
-///
-/// # Parameters
-///
-/// - `start`: The starting value of the range.
-/// - `end`: The end value of the range (inclusive).
-/// - `num`: The number of points to generate in the range.
-///
-/// # Returns
-///
-/// A vector containing the generated floating point values.
-#[must_use]
-pub fn float_vec_linspace(
- start: f64,
- end: f64,
- num: usize,
-) -> Vec {
- assert!(num > 0, "Number of points must be positive");
-
- if num == 1 {
- return vec![start];
- }
-
- let step = (end - start) / (num as f64 - 1.0);
-
- let end = if step > 0.0 {
- end + f64::EPSILON // Avoid floating point precision issues
- } else {
- end - f64::EPSILON // Avoid floating point precision issues
- };
-
- float_vec_arange(start, end, Some(step))
-}
-/// Create a 1D tensor with evenly spaced floating point values.
-///
-/// This function generates a tensor with values starting from `start`, ending at `end`, and incrementing by `step`.
-/// If `step` is not provided, it defaults to `1.0` if `start < end`, or `-1.0` if `start > end`.
-///
-/// # Parameters
-///
-/// - `start`: The starting value of the range.
-/// - `end`: The end value of the range (exclusive).
-/// - `step`: An optional step value. If not provided, defaults to `1.0` or `-1.0` based on the order of `start` and `end`.
-///
-/// # Returns
-///
-/// A 1D tensor containing the generated floating point values.
-#[must_use]
-pub fn float_arange(
- start: f64,
- end: f64,
- step: Option,
- device: &B::Device,
-) -> Tensor {
- let values = float_vec_arange(start, end, step);
- Tensor::from_data(values.as_slice(), device)
-}
-
-/// Create a 1D tensor with evenly spaced floating point values.
-///
-/// This function generates a tensor with `num` values starting from `start`, ending at `end`, and evenly spaced.
-///
-/// # Parameters
-///
-/// - `start`: The starting value of the range.
-/// - `end`: The end value of the range (inclusive).
-/// - `num`: The number of points to generate in the range.
-///
-/// # Returns
-///
-/// A 1D tensor containing the generated floating point values.
-#[must_use]
-pub fn float_linspace(
- start: f64,
- end: f64,
- num: usize,
- device: &B::Device,
-) -> Tensor {
- let values = float_vec_linspace(start, end, num);
- Tensor::from_data(values.as_slice(), device)
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use burn::backend::NdArray;
- use burn::prelude::TensorData;
-
- #[test]
- fn test_float_arange() {
- let device = Default::default();
- let start: f64 = 3.0;
- let end: f64 = -1.0 - f64::EPSILON;
-
- let actual = float_arange::(start, end, None, &device);
-
- actual
- .to_data()
- .assert_eq(&TensorData::from([3.0, 2.0, 1.0, 0.0, -1.0]), false);
- }
-
- #[should_panic(expected = "Step must be negative when start > end")]
- #[test]
- fn test_float_arange_panic_step_negative() {
- let device = Default::default();
- // This should panic because the step is not negative
- let _ = float_arange::(3.0, -1.0, Some(1.0), &device);
- }
-
- #[should_panic(expected = "Step must be positive when start < end")]
- #[test]
- fn test_float_arange_panic_step_positive() {
- let device = Default::default();
- // This should panic because the step is not positive
- let _ = float_arange::(-1.0, 3.0, Some(-1.0), &device);
- }
-
- #[test]
- fn test_float_vec_linspace_int_step() {
- let device = Default::default();
- let start: f64 = 0.0;
- let end: f64 = 1.0;
- let num: usize = 5;
-
- let actual = float_linspace::(start, end, num, &device);
-
- actual
- .to_data()
- .assert_eq(&TensorData::from([0.0, 0.25, 0.5, 0.75, 1.0]), false);
- }
-
- #[test]
- fn test_float_vec_linspace_neg_float_step() {
- let device = Default::default();
- let start: f64 = 1.0;
- let end: f64 = -0.2;
- let num: usize = 5;
-
- let actual = float_linspace::(start, end, num, &device);
-
- actual
- .to_data()
- .assert_eq(&TensorData::from([1.0, 0.7, 0.4, 0.1, -0.2]), false);
- }
-
- #[test]
- fn test_float_vec_linspace_n1() {
- let device = Default::default();
- let start: f64 = 0.0;
- let end: f64 = 1.0;
- let num: usize = 1;
-
- let actual = float_linspace::(start, end, num, &device);
- // println!("{actual:?}");
-
- actual.to_data().assert_eq(&TensorData::from([0.0]), false);
- }
-}
diff --git a/crates/bimm/src/compat/type_mapper.rs b/crates/bimm/src/compat/type_mapper.rs
new file mode 100644
index 00000000..d56527c1
--- /dev/null
+++ b/crates/bimm/src/compat/type_mapper.rs
@@ -0,0 +1,39 @@
+//! Type mapper for changing the data type of tensors in a module.
+use burn::{
+ Tensor,
+ module::{
+ ModuleMapper,
+ Param,
+ },
+ prelude::Backend,
+ tensor::DType,
+};
+
+/// Type mapper for changing the data type of tensors in a module.
+pub struct DTypeMapper {
+ /// Target data type for tensor conversion.
+ pub dt: DType,
+
+ /// Phantom data to ensure that `B` is a valid backend.
+ phantom: std::marker::PhantomData,
+}
+
+impl DTypeMapper {
+ /// Create a new type mapper with the specified target data type.
+ pub fn new(dt: DType) -> Self {
+ Self {
+ dt,
+ phantom: std::marker::PhantomData,
+ }
+ }
+}
+
+impl ModuleMapper for DTypeMapper {
+ fn map_float(
+ &mut self,
+ param: Param>,
+ ) -> Param> {
+ let (id, tensor, mapper) = param.consume();
+ Param::from_mapped_value(id, tensor.cast(self.dt), mapper)
+ }
+}
diff --git a/crates/bimm/src/layers/blocks/cna.rs b/crates/bimm/src/layers/blocks/cna.rs
deleted file mode 100644
index c0942e67..00000000
--- a/crates/bimm/src/layers/blocks/cna.rs
+++ /dev/null
@@ -1,360 +0,0 @@
-//! # `CNA2d` - conv/norm/activation block.
-//!
-//! A [`CNA2d`] module is:
-//! * a [`Conv2d`] layer,
-//! * a [`Normalization`] layer,
-//! * a [`Activation`] layer.
-//!
-//! With support for hooking the forward method,
-//! to run code between the norm and application layers.
-
-use bimm_contracts::{assert_shape_contract_periodically, unpack_shape_contract};
-use burn::config::Config;
-use burn::module::Module;
-use burn::nn::activation::{Activation, ActivationConfig};
-use burn::nn::conv::{Conv2d, Conv2dConfig};
-use burn::nn::norm::{Normalization, NormalizationConfig};
-use burn::prelude::{Backend, Tensor};
-
-/// Abstract policy for [`CNA2d`] Config.
-///
-/// Defines a [`NormalizationConfig`] and [`ActivationConfig`],
-/// and can be lifted to a [`CNA2dConfig`] to match a [`Conv2dConfig`].
-///
-/// The abstract [`NormalizationConfig`] will be feature matched
-/// with the target [`Conv2dConfig`].
-#[derive(Config, Debug)]
-pub struct AbstractCNA2dConfig {
- /// The [`Normalization`] config.
- pub norm: NormalizationConfig,
-
- /// Activation Config.
- #[config(default = "ActivationConfig::Relu")]
- pub act: ActivationConfig,
-}
-
-impl AbstractCNA2dConfig {
- /// Merge with a [`Conv2dConfig`] to construct a [`CNA2dConfig`].
- ///
- /// The abstract [`NormalizationConfig`] will be feature matched
- /// with the target [`Conv2dConfig`], resulting in a normalization
- /// layer sized appropriately for the input convolution.
- pub fn build_config(
- &self,
- conv: Conv2dConfig,
- ) -> CNA2dConfig {
- CNA2dConfig {
- conv,
- norm: self.norm.clone(),
- act: self.act.clone(),
- }
- .match_norm_features()
- }
-}
-
-/// [`CNA2d`] Meta.
-pub trait CNA2dMeta {
- /// Number of input channels.
- fn in_channels(&self) -> usize;
-
- /// Number of output channels.
- fn out_channels(&self) -> usize;
-
- /// Number of groups.
- fn groups(&self) -> usize;
-
- /// Get the stride.
- fn stride(&self) -> [usize; 2];
-}
-
-/// [`CNA2d`] Config.
-///
-/// Implements [`CNA2dMeta`].
-#[derive(Config, Debug)]
-pub struct CNA2dConfig {
- /// The [`Conv2d`] config.
- pub conv: Conv2dConfig,
-
- /// The [`Normalization`] config.
- pub norm: NormalizationConfig,
-
- /// The [`Activation`] config.
- #[config(default = "ActivationConfig::Relu")]
- pub act: ActivationConfig,
-}
-
-impl CNA2dMeta for CNA2dConfig {
- fn in_channels(&self) -> usize {
- self.conv.channels[0]
- }
-
- fn out_channels(&self) -> usize {
- self.conv.channels[1]
- }
-
- fn groups(&self) -> usize {
- self.conv.groups
- }
-
- fn stride(&self) -> [usize; 2] {
- self.conv.stride
- }
-}
-
-impl CNA2dConfig {
- /// Initialize a [`CNA2d`].
- ///
- /// Auto-matches the norm layer input channels
- /// to the conv layer's output channels.
- pub fn init(
- self,
- device: &B::Device,
- ) -> CNA2d {
- let out_channels = self.out_channels();
- CNA2d {
- conv: self.conv.init(device),
- norm: self.norm.with_num_features(out_channels).init(device),
- act: self.act.init(device),
- }
- }
-
- /// Adjust the norm features to match the conv output size.
- ///
- /// ['`CNA2dConfig::init`'] does this automatically.
- pub fn match_norm_features(self) -> Self {
- let features = self.out_channels();
- let norm = self.norm.with_num_features(features);
- Self { norm, ..self }
- }
-}
-
-/// Sequenced conv/norm/activation block.
-///
-/// A [`CNA2d`] module is:
-/// * a [`Conv2d`] layer,
-/// * a [`Normalization`] layer,
-/// * a [`Activation`] layer.
-///
-/// With support for hooking the forward method,
-/// to run code between the norm and application layers.
-///
-/// Implements [`CNA2dMeta`].
-#[derive(Module, Debug)]
-pub struct CNA2d {
- /// Internal Conv2d layer.
- pub conv: Conv2d,
-
- /// Internal Norm Layer.
- pub norm: Normalization,
-
- /// Activation layer.
- pub act: Activation,
-}
-
-impl CNA2dMeta for CNA2d {
- fn in_channels(&self) -> usize {
- self.conv.weight.dims()[1] * self.groups()
- }
-
- fn out_channels(&self) -> usize {
- self.conv.weight.dims()[0]
- }
-
- fn groups(&self) -> usize {
- self.conv.groups
- }
-
- fn stride(&self) -> [usize; 2] {
- self.conv.stride
- }
-}
-
-impl CNA2d {
- /// Forward Pass.
- ///
- /// Applies the conv/norm/act layers in sequence:
- ///
- /// ```rust,ignore
- /// let x = self.conv.forward(input);
- /// let x = self.norm.forward(x);
- /// let x = self.act.forward(x);
- /// return x
- /// ```
- ///
- /// # Arguments
- ///
- /// - `input`: \
- /// ``[batch, in_channels, in_height=out_height*stride, in_width=out_width*stride]``.
- ///
- /// # Returns
- ///
- /// ``[batch, out_channels, out_height, out_width]``
- pub fn forward(
- &self,
- input: Tensor,
- ) -> Tensor {
- self.map_forward(input, |x| x)
- }
-
- /// Mapping Forward Pass.
- ///
- /// Applies the callback fn after normalization but before activation.
- ///
- /// ```rust,ignore
- /// let x = self.conv.forward(input);
- /// let x = self.norm.forward(x);
- /// let x = f(x);
- /// let x = self.act.forward(x);
- /// return x
- /// ```
- ///
- /// # Arguments
- ///
- /// - `input`: \
- /// ``[batch, in_channels, in_height=out_height*stride, in_width=out_width*stride]``.
- /// - `f`: a callback endofunction, from/to ``[batch, in_channels, out_height, out_width]``.
- ///
- /// # Returns
- ///
- /// ``[batch, out_channels, out_height, out_width]``
- pub fn map_forward(
- &self,
- input: Tensor,
- f: F,
- ) -> Tensor
- where
- F: FnOnce(Tensor) -> Tensor,
- {
- let [batch, out_height, out_width] = unpack_shape_contract!(
- [
- "batch",
- "in_channels",
- "in_height" = "out_height" * "height_stride",
- "in_width" = "out_width" * "width_stride"
- ],
- &input.dims(),
- &["batch", "out_height", "out_width"],
- &[
- ("in_channels", self.in_channels()),
- ("height_stride", self.stride()[0]),
- ("width_stride", self.stride()[1]),
- ]
- );
- let x = self.conv.forward(input);
-
- assert_shape_contract_periodically!(
- ["batch", "out_channels", "out_height", "out_width"],
- &x.dims(),
- &[
- ("batch", batch),
- ("out_channels", self.out_channels()),
- ("out_height", out_height),
- ("out_width", out_width)
- ]
- );
-
- let x = self.norm.forward(x);
-
- let x = f(x);
-
- let x = self.act.forward(x);
-
- assert_shape_contract_periodically!(
- ["batch", "out_channels", "out_height", "out_width"],
- &x.dims(),
- &[
- ("batch", batch),
- ("out_channels", self.out_channels()),
- ("out_height", out_height),
- ("out_width", out_width)
- ]
- );
-
- x
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use burn::backend::{Autodiff, NdArray};
- use burn::nn::activation::ActivationConfig;
- use burn::nn::norm::NormalizationConfig;
- use burn::nn::{BatchNormConfig, PaddingConfig2d};
- use burn::tensor::Distribution;
-
- #[test]
- fn test_conv_norm_config() {
- let abstract_config =
- AbstractCNA2dConfig::new(NormalizationConfig::Batch(BatchNormConfig::new(0)));
-
- let conv_config = Conv2dConfig::new([2, 4], [3, 3])
- .with_stride([2, 2])
- .with_padding(PaddingConfig2d::Explicit(1, 1))
- .with_bias(false);
-
- let config: CNA2dConfig = abstract_config.build_config(conv_config.clone());
-
- assert_eq!(config.in_channels(), 2);
- assert_eq!(config.out_channels(), 4);
- assert_eq!(config.groups(), 1);
- assert_eq!(config.stride(), [2, 2]);
- }
-
- #[test]
- fn test_cna() {
- type B = Autodiff>;
- let device = Default::default();
-
- let config = CNA2dConfig::new(
- Conv2dConfig::new([2, 4], [3, 3])
- .with_stride([2, 2])
- .with_padding(PaddingConfig2d::Explicit(1, 1))
- .with_bias(false),
- NormalizationConfig::Batch(BatchNormConfig::new(0)),
- )
- .with_act(ActivationConfig::Relu);
-
- let layer: CNA2d = config.init(&device);
- assert_eq!(layer.in_channels(), 2);
- assert_eq!(layer.out_channels(), 4);
- assert_eq!(layer.groups(), 1);
- assert_eq!(layer.stride(), [2, 2]);
-
- let batch_size = 2;
- let height = 10;
- let width = 10;
- let channels = 2;
-
- let input = Tensor::random(
- [batch_size, channels, height, width],
- Distribution::Default,
- &device,
- );
-
- {
- let output = layer.forward(input.clone());
- let expected = {
- let x = layer.conv.forward(input.clone());
- let x = layer.norm.forward(x);
- let x = layer.act.forward(x);
- x
- };
- output.to_data().assert_eq(&expected.to_data(), true);
- }
-
- {
- let hook = |x| x * 2.0;
-
- let output = layer.map_forward(input.clone(), hook);
- let expected = {
- let x = layer.conv.forward(input.clone());
- let x = layer.norm.forward(x);
- let x = hook(x);
- let x = layer.act.forward(x);
- x
- };
- output.to_data().assert_eq(&expected.to_data(), true);
- }
- }
-}
diff --git a/crates/bimm/src/layers/blocks/conv_norm.rs b/crates/bimm/src/layers/blocks/conv_norm.rs
deleted file mode 100644
index bd15d9fa..00000000
--- a/crates/bimm/src/layers/blocks/conv_norm.rs
+++ /dev/null
@@ -1,176 +0,0 @@
-//! # `ConvNorm` Module
-//!
-//! A [`ConvNorm2d`] module is a [`Conv2d`] layer followed by a [`BatchNorm`] layer.
-
-use bimm_contracts::{assert_shape_contract_periodically, unpack_shape_contract};
-use burn::config::Config;
-use burn::module::Module;
-use burn::nn::conv::{Conv2d, Conv2dConfig};
-use burn::nn::{BatchNorm, BatchNormConfig, Initializer};
-use burn::prelude::{Backend, Tensor};
-
-/// [`ConvNorm2d`] Meta.
-pub trait ConvNorm2dMeta {
- /// Number of input channels.
- fn in_channels(&self) -> usize;
-
- /// Number of groups.
- fn groups(&self) -> usize;
-
- /// Number of output channels.
- fn out_channels(&self) -> usize;
-
- /// Get the stride.
- fn stride(&self) -> &[usize; 2];
-}
-
-/// [`ConvNorm2d`] Config.
-#[derive(Config, Debug)]
-pub struct ConvNorm2dConfig {
- /// The [`Conv2d`] config.
- pub conv: Conv2dConfig,
-}
-
-impl ConvNorm2dMeta for ConvNorm2dConfig {
- fn in_channels(&self) -> usize {
- self.conv.channels[0]
- }
-
- fn groups(&self) -> usize {
- self.conv.groups
- }
-
- fn out_channels(&self) -> usize {
- self.conv.channels[1]
- }
-
- fn stride(&self) -> &[usize; 2] {
- &self.conv.stride
- }
-}
-
-impl From for ConvNorm2dConfig {
- fn from(conv: Conv2dConfig) -> Self {
- Self { conv }
- }
-}
-
-impl ConvNorm2dConfig {
- /// Initialize a [`ConvNorm2d`].
- pub fn init(
- self,
- device: &B::Device,
- ) -> ConvNorm2d {
- ConvNorm2d {
- conv: self.conv.init(device),
-
- norm: BatchNormConfig::new(self.conv.channels[1]).init(device),
- }
- }
-
- /// Set the initializer.
- pub fn with_initializer(
- self,
- initializer: Initializer,
- ) -> Self {
- Self {
- conv: self.conv.with_initializer(initializer),
- }
- }
-}
-
-/// Grouped [`Conv2d`] and [`BatchNorm`] layer.
-#[derive(Module, Debug)]
-pub struct ConvNorm2d {
- /// Internal Conv2d layer.
- pub conv: Conv2d,
-
- /// Internal Norm Layer.
- pub norm: BatchNorm,
-}
-
-impl ConvNorm2dMeta for ConvNorm2d {
- fn in_channels(&self) -> usize {
- self.conv.weight.shape().dims[1] * self.groups()
- }
-
- fn groups(&self) -> usize {
- self.conv.groups
- }
-
- fn out_channels(&self) -> usize {
- self.conv.weight.shape().dims[0]
- }
-
- fn stride(&self) -> &[usize; 2] {
- &self.conv.stride
- }
-}
-
-impl ConvNorm2d {
- /// Zero initialize the norm layer's weights.
- ///
- /// This is used by / referenced in upstream `ResNet` init.
- ///
- /// TODO: Track down the paper that recommends this.
- pub fn zero_init_norm(&mut self) {
- self.norm.gamma = self.norm.gamma.clone().map(|p| p.slice_fill([..], 0.0));
- }
- /// Forward Pass.
- pub fn forward(
- &self,
- input: Tensor,
- ) -> Tensor {
- let [batch, out_height, out_width] = unpack_shape_contract!(
- [
- "batch",
- "in_channels",
- "in_height" = "out_height" * "height_stride",
- "in_width" = "out_width" * "width_stride"
- ],
- &input.dims(),
- &["batch", "out_height", "out_width"],
- &[
- ("in_channels", self.in_channels()),
- ("height_stride", self.stride()[0]),
- ("width_stride", self.stride()[1]),
- ]
- );
- let x = self.conv.forward(input);
-
- let x = self.norm.forward(x);
-
- assert_shape_contract_periodically!(
- ["batch", "out_channels", "out_height", "out_width"],
- &x.dims(),
- &[
- ("batch", batch),
- ("out_channels", self.out_channels()),
- ("out_height", out_height),
- ("out_width", out_width)
- ]
- );
-
- x
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use burn::nn::PaddingConfig2d;
-
- #[test]
- fn test_conv_norm_config() {
- let inner_config = Conv2dConfig::new([2, 4], [3, 3])
- .with_stride([2, 2])
- .with_padding(PaddingConfig2d::Explicit(1, 1))
- .with_bias(false);
-
- let config: ConvNorm2dConfig = inner_config.clone().into();
-
- assert_eq!(&config.conv.channels, &inner_config.channels);
- assert_eq!(&config.conv.kernel_size, &inner_config.kernel_size);
- assert_eq!(&config.conv.stride, &inner_config.stride);
- }
-}
diff --git a/crates/bimm/src/layers/blocks/mod.rs b/crates/bimm/src/layers/blocks/mod.rs
deleted file mode 100644
index 9a9ab529..00000000
--- a/crates/bimm/src/layers/blocks/mod.rs
+++ /dev/null
@@ -1,5 +0,0 @@
-//! # Common Building Blocks
-
-pub mod conv_norm;
-
-pub mod cna;
diff --git a/crates/bimm/src/layers/drop/drop_block.rs b/crates/bimm/src/layers/drop/drop_block.rs
deleted file mode 100644
index e3a14ba3..00000000
--- a/crates/bimm/src/layers/drop/drop_block.rs
+++ /dev/null
@@ -1,733 +0,0 @@
-//! # `DropBlock` Layers
-//!
-//! Based upon [DropBlock (Ghiasi et al., 2018)](https://arxiv.org/pdf/1810.12890.pdf);
-//! inspired also by the `python-image-models` implementation.
-
-use crate::layers::drop::size_config::SizeConfig;
-use crate::utility::burn::kernels;
-use crate::utility::burn::noise::NoiseConfig;
-use crate::utility::probability::expect_probability;
-use alloc::format;
-use bimm_contracts::unpack_shape_contract;
-use burn::config::Config;
-use burn::module::{Content, Module, ModuleDisplay, ModuleDisplayDefault};
-use burn::prelude::{Backend, Float, Tensor};
-use burn::tensor::module::max_pool2d;
-use burn::tensor::{DType, Distribution};
-use core::ops::Range;
-use serde::{Deserialize, Serialize};
-
-/// Configuration for `DropBlock`.
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
-pub struct DropBlockOptions {
- /// The drop probability.
- pub drop_prob: f64,
-
- /// The block size.
- pub kernel: [SizeConfig; 2],
-
- /// The gamma scale.
- pub gamma_scale: f64,
-
- /// Whether to compute batchwise blocks selection noise.
- /// The alternative is to compute per-image/channel noise.
- /// This results in a speedup proportional to the batchsize.
- pub batchwise: bool,
-
- /// Should color drop be coupled, or independent?
- pub couple_channels: bool,
-
- /// Permit partial blocks at the edges.
- /// This results in a significant speedup when using noise.
- pub partial_edge_blocks: bool,
-
- /// The noise configuration.
- pub noise: Option,
-}
-
-impl ModuleDisplay for DropBlockOptions {}
-impl ModuleDisplayDefault for DropBlockOptions {
- fn content(
- &self,
- content: Content,
- ) -> Option {
- Some(
- content
- .add("drop_prob", &self.drop_prob)
- .add("kernel", format!("{:?}", &self.kernel).as_str())
- .add("gamma_scale", &self.gamma_scale)
- .add("batchwise", &self.batchwise)
- .add("couple_channels", &self.couple_channels)
- .add("partial_edge_blocks", &self.partial_edge_blocks)
- .add("noise_cfg", &self.noise),
- )
- }
-}
-
-impl Default for DropBlockOptions {
- fn default() -> Self {
- Self {
- drop_prob: 0.1,
- kernel: [7.into(); 2],
- gamma_scale: 1.0,
- noise: None,
- batchwise: true,
- couple_channels: false,
- partial_edge_blocks: false,
- }
- }
-}
-
-impl DropBlockOptions {
- /// Extend the options with the given probability.
- ///
- /// # Arguments
- ///
- /// - `drop_prob` - the probability.
- ///
- /// # Panics
- ///
- /// If the `drop_prob` is not in ``[0.0, 1.0]``.
- pub fn with_drop_prob(
- self,
- drop_prob: f64,
- ) -> Self {
- Self {
- drop_prob: expect_probability(drop_prob),
- ..self
- }
- }
-
- /// Set the symmetric kernel block size.
- ///
- /// This is equivalent to `.with_kernel([block_size; 2])`
- ///
- /// # Arguments
- ///
- /// - `block_size` - the symmetric kernel size.
- pub fn with_block_size(
- self,
- block_size: S,
- ) -> Self
- where
- S: Into,
- {
- Self {
- kernel: [block_size.into(); 2],
- ..self
- }
- }
-
- /// Sets the kernel size.
- ///
- /// # Arguments
- ///
- /// - `kernel` - the kernel size.
- pub fn with_kernel(
- self,
- kernel: [S; 2],
- ) -> Self
- where
- S: Into,
- {
- let [x, y] = kernel;
- let kernel = [x.into(), y.into()];
- Self { kernel, ..self }
- }
-
- /// Set the gamma scale.
- pub fn with_gamma_scale(
- self,
- gamma_scale: f64,
- ) -> Self {
- Self {
- gamma_scale,
- ..self
- }
- }
-
- /// Set whether to use noise.
- pub fn with_noise(
- self,
- noise_cfg: N,
- ) -> Self
- where
- N: Into>,
- {
- Self {
- noise: noise_cfg.into(),
- ..self
- }
- }
-
- /// Set if batchwise noise should be used.
- ///
- /// When `batchwise` is enabled, the entire batch shares the same noise;
- /// otherwise, each image gets its own noise.
- ///
- /// Depending on the batchsize, this could be a significant speedup.
- ///
- /// # Arguments
- ///
- /// - `batchwise` - the mode.
- pub fn with_batchwise(
- self,
- batchwise: bool,
- ) -> Self {
- Self { batchwise, ..self }
- }
-
- /// Set if channels should be coupled.
- pub fn with_couple_channels(
- self,
- couple_channels: bool,
- ) -> Self {
- Self {
- couple_channels,
- ..self
- }
- }
-
- /// Set if partial edge blocks should be used.
- ///
- /// Partial edge blocks are blocks which overlap the edge;
- /// and are smaller than the full kernel; but not performing
- /// the edge-check saves compute time.
- ///
- /// # Arguments
- ///
- /// - `partial_edge_blocks`: the mode.
- pub fn with_partial_edge_blocks(
- self,
- partial_edge_blocks: bool,
- ) -> Self {
- Self {
- partial_edge_blocks,
- ..self
- }
- }
-
- /// Clip the kernel to fit within the shape.
- ///
- /// # Arguments
- ///
- /// - `shape`: the shape of the target image.
- ///
- /// # Returns
- ///
- /// The dimension-wise clipped kernel shape.
- #[inline]
- pub fn clipped_kernel(
- &self,
- shape: [usize; 2],
- ) -> [usize; 2] {
- let [h, w] = shape;
- let [kh, kw] = self.kernel;
- let kh = kh.resolve(h);
- let kw = kw.resolve(w);
- [core::cmp::min(h, kh), core::cmp::min(w, kw)]
- }
-
- /// Compute the adjusted gamma rate.
- ///
- /// Gamma is the adjusted probability that any given point is the midpoint
- /// of a dropped block; given the desired `drop_rate`, the block size, and the input size.
- ///
- /// # Arguments
- ///
- /// - `shape`: the shape of the target tensor.
- #[inline]
- pub fn gamma(
- &self,
- shape: [usize; 2],
- ) -> f64 {
- let [h, w] = shape;
- let [kh, kw] = self.clipped_kernel(shape);
-
- (self.gamma_scale * self.drop_prob * ((h * w) as f64))
- / ((kh * kw) as f64)
- / (((h - kh + 1) * (w - kw + 1)) as f64)
- }
-
- /// Compute the gamma noise for the given noise batch shape.
- ///
- /// # Args
- ///
- /// - `shape`: the target noise batch shape.
- ///
- /// # Returns
- ///
- /// Gamma noise, sampled at ``self.gamma([h, w])`` rate.
- pub fn gamma_noise(
- &self,
- noise_shape: [usize; 4],
- device: &B::Device,
- ) -> Tensor {
- let [_, _, h, w] = noise_shape;
- let gamma = self.gamma([h, w]);
- Tensor::random(noise_shape, Distribution::Bernoulli(gamma), device)
- }
-}
-
-/// Convert drop block gamma noise to a selection filter.
-///
-/// This is a deterministic internal component of `drop_block`.
-///
-/// # Arguments
-///
-/// * `selected_blocks` - Input selection noise;
-/// `1.0` at the midpoints of selected blocks to drop,
-/// `0.0` everywhere else. Expected to be gamma noise.
-/// * `kernel_shape` - the shape of the kernel.
-/// * `partial_edge_blocks` - permit partial blocks at the edges, faster.
-fn drop_block_2d_drop_filter_(
- selected_blocks: Tensor,
- kernel_shape: [usize; 2],
- partial_edge_blocks: bool,
-) -> Tensor {
- let [_, _, h, w] = unpack_shape_contract!(["b", "c", "h", "w"], &selected_blocks.dims());
- let [kh, kw] = kernel_shape;
-
- assert!(
- kh <= h && kw <= w,
- "Kernel size ({kh}, {kw}) is larger than input size ({h}, {w})",
- );
-
- let dtype = selected_blocks.dtype();
- let device = &selected_blocks.device();
-
- let mut selection = selected_blocks;
-
- if !partial_edge_blocks {
- selection = selection
- * kernels::conv2d_kernel_midpoint_filter::([h, w], kernel_shape, device)
- .unsqueeze_dims::<4>(&[0, 1])
- .cast(dtype);
- }
-
- let ceil_mode = false;
-
- selection = max_pool2d(
- selection,
- kernel_shape,
- [1, 1],
- [kh / 2, kw / 2],
- [1, 1],
- ceil_mode,
- );
-
- // Clip even-kernel padding artifacts.
- if (kh % 2) == 0 || (kw % 2) == 0 {
- let mut ranges: [Range; 4] = selection.shape().into_ranges().try_into().unwrap();
- ranges[2].start = ((kh % 2) == 0) as usize;
- ranges[3].start = ((kw % 2) == 0) as usize;
-
- selection = selection.slice(ranges);
- }
-
- selection
-}
-
-/// `drop_block_2d`
-///
-/// Drops block kernels from a tensor; many options.
-///
-/// Dropped values can be resampled from a noise distribution,
-/// kept values can be re-normalized.
-/// The drop probability, block size, and several performance/quality tradeoffs
-/// can be configured.
-///
-/// Based upon [DropBlock (Ghiasi, et al., 2018)](https://arxiv.org/pdf/1810.12890.pdf);
-/// inspired also by the `python-image-models` implementation.
-///
-/// # Arguments
-///
-/// * `tensor` - the tensor to operate on, ``[batch, channels, height, width]``.
-/// * `options` - the algorithm options.
-///
-/// # Returns
-///
-/// A ``[batch, channels, height, width]`` tensor.
-pub fn drop_block_2d(
- tensor: Tensor,
- options: &DropBlockOptions,
-) -> Tensor {
- if options.drop_prob == 0.0 {
- // This is a no-op.
- return tensor;
- }
-
- let [b, c, h, w] = tensor.shape().dims();
- let kernel = options.clipped_kernel([h, w]);
-
- let t_shape = tensor.shape();
- let device = &tensor.device();
- let dtype = tensor.dtype();
-
- let noise_shape = [
- if options.batchwise { 1 } else { b },
- if options.couple_channels { 1 } else { c },
- h,
- w,
- ];
-
- let gamma_noise = options.gamma_noise(noise_shape, device);
-
- let drop_filter: Tensor =
- drop_block_2d_drop_filter_(gamma_noise, kernel, options.partial_edge_blocks).cast(dtype);
- let keep_filter: Tensor = 1.0 - drop_filter.clone();
-
- if let Some(noise_cfg) = &options.noise {
- // Fill in the dropped regions with sampled noise.
- let noise: Tensor = noise_cfg.noise(noise_shape, device).cast(dtype);
- let noise = noise * drop_filter;
-
- tensor * keep_filter.expand(t_shape.clone()) + noise.expand(t_shape)
- } else {
- // Rescale to normalize to 1.0.
- let count = keep_filter.shape().num_elements() as f32;
- let total = keep_filter.clone().cast(DType::F32).sum();
- let norm_scale = count / total.add_scalar(1e-7);
-
- tensor * keep_filter.expand(t_shape.clone()) * norm_scale.cast(dtype).expand(t_shape)
- }
-}
-
-/// Config for [`DropBlock2dConfig`] modules.
-#[derive(Config, Debug)]
-pub struct DropBlock2dConfig {
- /// The options for the drop block algorithm.
- #[config(default = "DropBlockOptions::default()")]
- pub options: DropBlockOptions,
-}
-
-impl From for DropBlock2dConfig {
- fn from(options: DropBlockOptions) -> Self {
- Self { options }
- }
-}
-
-impl DropBlock2dConfig {
- /// Initialize a [`DropBlock2d`] module.
- pub fn init(&self) -> DropBlock2d {
- DropBlock2d {
- options: self.options.clone(),
- }
- }
-}
-
-/// `DropBlock2d`
-///
-/// A module that applies drop block (when training).
-///
-/// Based upon [DropBlock (Ghiasi et al., 2018)](https://arxiv.org/pdf/1810.12890.pdf);
-/// inspired also by the `python-image-models` implementation.
-#[derive(Module, Clone, Debug)]
-pub struct DropBlock2d {
- /// The options for the drop block algorithm.
- pub options: DropBlockOptions,
-}
-
-impl DropBlock2d {
- /// When training, applies `drop_block_2d` to the tensor;
- /// otherwise, is a no-op pass-through.
- ///
- /// # Arguments
- ///
- /// - `tensor` - the input.
- ///
- /// # Returns
- ///
- /// A tensor of the same shape, type, and device.
- pub fn forward(
- &self,
- tensor: Tensor,
- ) -> Tensor {
- if B::ad_enabled() {
- drop_block_2d(tensor.clone(), &self.options)
- } else {
- tensor
- }
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use crate::utility::burn::noise::NoiseConfig;
- use burn::backend::{Autodiff, NdArray};
- use burn::module::DisplaySettings;
- use burn::prelude::TensorData;
-
- #[test]
- fn test_drop_block_options() {
- let options = DropBlockOptions::default();
- assert_eq!(options.drop_prob, 0.1);
- assert_eq!(options.kernel, [SizeConfig::Fixed(7); 2]);
- assert_eq!(options.gamma_scale, 1.0);
- assert_eq!(options.couple_channels, false);
- assert!(options.noise.is_none());
- assert_eq!(options.batchwise, true);
-
- let options = options.with_couple_channels(true);
- assert_eq!(options.couple_channels, true);
-
- let options = options.with_drop_prob(0.2);
- assert_eq!(options.drop_prob, 0.2);
-
- let options = options.with_block_size(0.25);
- assert_eq!(options.kernel, [SizeConfig::Ratio(0.25); 2]);
-
- let options = options.with_gamma_scale(0.5);
- assert_eq!(options.gamma_scale, 0.5);
-
- let options = options.with_batchwise(false);
- assert_eq!(options.batchwise, false);
- }
-
- #[test]
- fn test_drop_block_options_display() {
- let options = DropBlockOptions::default().with_noise(NoiseConfig::default());
- let settings = DisplaySettings::default();
-
- assert_eq!(
- options.format(settings),
- indoc::indoc! {r#"
- DropBlockOptions {
- drop_prob: 0.1
- kernel: [Fixed(7), Fixed(7)]
- gamma_scale: 1
- batchwise: true
- couple_channels: false
- partial_edge_blocks: false
- noise_cfg: NoiseConfig {
- distribution: Distribution::Default
- clamp: None
- }
- }"#
- }
- )
- }
-
- #[test]
- fn test_gamma() {
- let options = DropBlockOptions::default()
- .with_drop_prob(0.1)
- .with_gamma_scale(1.2)
- .with_kernel([2, 3]);
-
- let shape = [7, 9];
- let [h, w] = shape;
- let [kh, kw] = options.clipped_kernel(shape);
-
- let total_size = (h * w) as f64;
-
- let gamma = options.gamma([h, w]);
-
- let expected = (options.gamma_scale * options.drop_prob * total_size)
- / ((kh * kw) as f64)
- / (((h - kh + 1) * (w - kw + 1)) as f64);
-
- assert_eq!(gamma, expected);
- }
-
- #[test]
- fn test_drop_block_2d_drop_filter() {
- type B = NdArray;
- let device = Default::default();
-
- let selected_blocks: Tensor = Tensor::::from_data(
- [
- [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- ],
- &device,
- )
- .unsqueeze_dims::<4>(&[0, 1]);
-
- // No partial edge blocks.
- let drop_filter = drop_block_2d_drop_filter_(selected_blocks.clone(), [2, 3], false);
- drop_filter.squeeze_dims::<2>(&[0, 1]).to_data().assert_eq(
- &TensorData::from([
- [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0],
- [0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
- [0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- ]),
- false,
- );
-
- // Partial edge blocks.
- let drop_filter = drop_block_2d_drop_filter_(selected_blocks.clone(), [2, 3], true);
- drop_filter.squeeze_dims::<2>(&[0, 1]).to_data().assert_eq(
- &TensorData::from([
- [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0],
- [0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0],
- [0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0],
- [0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0],
- [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
- ]),
- false,
- );
- }
-
- #[test]
- fn test_drop_block_2d_no_op() {
- type B = NdArray;
- let device = Default::default();
-
- let shape = [2, 3, 7, 9];
- let tensor: Tensor = Tensor::ones(shape, &device);
-
- let drop_prob = 0.0;
-
- let drop = drop_block_2d(
- tensor.clone(),
- &DropBlockOptions::default()
- .with_drop_prob(drop_prob)
- .with_kernel([2, 3]),
- );
-
- drop.to_data().assert_eq(&tensor.to_data(), false);
- }
-
- #[test]
- fn test_drop_block_2d_with_norm() {
- type B = NdArray;
- let device = Default::default();
-
- let shape = [2, 3, 100, 100];
- let tensor: Tensor = Tensor::ones(shape, &device);
-
- let drop_prob = 0.1;
-
- let drop = drop_block_2d(
- tensor,
- &DropBlockOptions::default()
- .with_partial_edge_blocks(false)
- .with_drop_prob(drop_prob)
- .with_kernel([2, 3]),
- );
-
- // Count all 1.0; which are the non-dropped values.
- let numel = drop.shape().num_elements();
-
- // They've all been rescaled upwards.
- let keep_count = drop.clone().greater_elem(1.0).int().sum().into_scalar() as usize;
- let drop_count = numel - keep_count;
- let drop_ratio = drop_count as f64 / numel as f64;
- assert!((drop_ratio - drop_prob).abs() < 0.15);
-
- let total = drop.sum().into_scalar() as f64;
- let norm = total / numel as f64;
- assert!((norm - 1.0).abs() < 0.01);
- }
-
- #[test]
- fn test_drop_block_2d_with_noise() {
- type B = NdArray;
- let device = Default::default();
-
- let shape = [2, 3, 100, 100];
- let tensor: Tensor = Tensor::ones(shape, &device);
-
- let drop_prob = 0.1;
-
- let drop = drop_block_2d(
- tensor,
- &DropBlockOptions::default()
- .with_noise(NoiseConfig::default())
- .with_partial_edge_blocks(false)
- .with_drop_prob(drop_prob)
- .with_kernel([2, 3]),
- );
-
- // Count all 1.0; which are the non-dropped values.
- let numel = drop.shape().num_elements();
-
- // This should be an exact match; the Distribution::Default is [0.0, 1.0);
- // and will never generate a 1.0.
- let keep_count = drop.equal_elem(1.0).int().sum().into_scalar() as usize;
-
- let drop_count = numel - keep_count;
-
- let drop_ratio = drop_count as f64 / numel as f64;
-
- assert!((drop_ratio - drop_prob).abs() < 0.15);
- }
-
- #[test]
- fn test_module_inference() {
- type B = NdArray;
- let device = Default::default();
-
- let config = DropBlock2dConfig::new();
-
- let module = config.init();
-
- let batch_size = 2;
- let channels = 3;
- let height = 100;
- let width = height;
- let shape = [batch_size, channels, height, width];
-
- let tensor: Tensor = Tensor::ones(shape, &device);
-
- assert_eq!(B::ad_enabled(), false);
- let result = module.forward(tensor.clone());
-
- // Not under training; so a no-op.
- result.to_data().assert_eq(&tensor.to_data(), false);
- }
-
- #[test]
- fn test_module_training() {
- type B = Autodiff;
- let device = Default::default();
-
- let drop_prob = 0.1;
-
- let config = DropBlock2dConfig::new().with_options(
- DropBlockOptions::default()
- .with_drop_prob(drop_prob)
- .with_kernel([2, 3]),
- );
-
- let module = config.init();
-
- let batch_size = 2;
- let channels = 3;
- let height = 100;
- let width = height;
- let shape = [batch_size, channels, height, width];
-
- let tensor: Tensor = Tensor::ones(shape, &device);
-
- assert_eq!(B::ad_enabled(), true);
- let drop = module.forward(tensor.clone());
-
- // Count all 1.0; which are the non-dropped values.
- let numel = drop.shape().num_elements();
-
- // They've all been rescaled upwards.
- let keep_count = drop.clone().greater_elem(1.0).int().sum().into_scalar() as usize;
- let drop_count = numel - keep_count;
- let drop_ratio = drop_count as f64 / numel as f64;
- assert!((drop_ratio - drop_prob).abs() < 0.15);
-
- let total = drop.sum().into_scalar() as f64;
- let norm = total / numel as f64;
- assert!((norm - 1.0).abs() < 0.01);
- }
-}
diff --git a/crates/bimm/src/layers/drop/drop_path.rs b/crates/bimm/src/layers/drop/drop_path.rs
deleted file mode 100644
index fd7d4935..00000000
--- a/crates/bimm/src/layers/drop/drop_path.rs
+++ /dev/null
@@ -1,355 +0,0 @@
-//! Burn implementation of the `DropPath` (Stochastic Depth) regularization layer.
-//!
-//! Papers:
-//! `DropBlock`: A regularization method for convolutional networks ()
-//!
-//! Deep Networks with Stochastic Depth ()
-//!
-//! Inspired by the python implementation from the timm library:
-//!
-
-use crate::utility::probability;
-use burn::config::Config;
-use burn::module::Module;
-use burn::prelude::{Backend, Tensor};
-use burn::tensor::Distribution;
-
-/// `DropPath` (stochastic depth) regularization.
-///
-/// # Arguments
-///
-/// * `x`: Input tensor.
-/// * `drop_prob`: Probability of dropping a path.
-/// * `training`: Whether the model is in training mode.
-/// * `scale_by_keep`: Whether to scale the output by `1 / (1 - drop_prob)`
-///
-/// # Returns
-///
-/// * Output tensor with the same shape as the input tensor.
-#[must_use]
-pub fn drop_path(
- x: Tensor,
- drop_prob: f64,
- training: bool,
- scale_by_keep: bool,
-) -> Tensor {
- _drop_path_sample(
- x,
- drop_prob,
- training,
- scale_by_keep,
- |shape, keep_prob, device| {
- Tensor::::random(shape, Distribution::Bernoulli(keep_prob), device)
- },
- )
-}
-
-/// Internal implementation of `DropPath`.
-///
-/// Deferred to a separate function to allow for testing sampling.
-///
-/// # Arguments
-///
-/// * `x`: Input tensor.
-/// * `drop_prob`: Probability of dropping a path.
-/// * `training`: Whether the model is in training mode.
-/// * `scale_by_keep`: Whether to scale the output by `1 / (1 - drop_prob)`
-/// * `sample`: Sampling function to generate the random tensor.
-///
-/// # Returns
-///
-/// * Output tensor with the same shape as the input tensor.
-#[inline(always)]
-#[must_use]
-fn _drop_path_sample(
- x: Tensor,
- drop_prob: f64,
- training: bool,
- scale_by_keep: bool,
- sample: fn([usize; D], f64, &B::Device) -> Tensor,
-) -> Tensor {
- probability::expect_probability(drop_prob);
-
- if !training || drop_prob == 0.0 {
- return x;
- }
-
- let keep_prob = 1.0 - drop_prob;
-
- let mut shape = [1; D];
- shape[0] = x.dims()[0];
-
- let random_tensor = sample(shape, keep_prob, &x.device());
-
- let random_tensor = if keep_prob > 0.0 && scale_by_keep {
- random_tensor.div_scalar(keep_prob)
- } else {
- random_tensor
- };
-
- x * random_tensor
-}
-
-/// Common introspection interface for `DropPath` modules.
-pub trait DropPathMeta {
- /// Returns the drop probability.
- fn drop_prob(&self) -> f64;
-
- /// Returns the keep probability, which is `1.0 - drop_prob`.
- fn keep_prob(&self) -> f64 {
- 1.0 - self.drop_prob()
- }
-
- /// Returns whether the output is scaled by `1 / (1 - drop_prob)`.
- fn scale_by_keep(&self) -> bool;
-}
-
-/// Configuration for the `DropPath` module.
-#[derive(Config, Debug)]
-pub struct DropPathConfig {
- /// Probability of dropping a path.
- #[config(default = 0.0)]
- pub drop_prob: f64,
-
- /// Whether to scale the output by `1 / (1 - drop_prob)`.
- #[config(default = true)]
- pub scale_by_keep: bool,
-}
-
-impl DropPathMeta for DropPathConfig {
- fn drop_prob(&self) -> f64 {
- self.drop_prob
- }
-
- fn scale_by_keep(&self) -> bool {
- self.scale_by_keep
- }
-}
-
-impl DropPathConfig {
- /// Initializes a new `DropPath` module.
- #[inline(always)]
- #[must_use]
- pub fn init(&self) -> DropPath {
- DropPath {
- drop_prob: probability::expect_probability(self.drop_prob),
- scale_by_keep: self.scale_by_keep,
- }
- }
-}
-
-/// The `DropPath` module.
-///
-/// Burn Module that implements the `DropPath` (Stochastic Depth) regularization.
-#[derive(Module, Clone, Debug)]
-pub struct DropPath {
- /// Probability of dropping a path.
- pub drop_prob: f64,
-
- /// Whether to scale the output by `1 / (1 - drop_prob)`.
- pub scale_by_keep: bool,
-}
-
-impl DropPathMeta for DropPath {
- fn drop_prob(&self) -> f64 {
- self.drop_prob
- }
-
- fn scale_by_keep(&self) -> bool {
- self.scale_by_keep
- }
-}
-
-impl DropPath {
- /// Applies `drop_path` pass on the input tensor.
- #[must_use]
- pub fn forward(
- &self,
- input: Tensor,
- ) -> Tensor {
- let training = B::ad_enabled();
- drop_path(input, self.drop_prob, training, self.scale_by_keep)
- }
-
- /// Applies an inner function under conditional stochastic residual/depth-skip connection.
- ///
- /// This is used for stochastic depth in the transformer block.
- ///
- /// # Parameters
- ///
- /// * `x` - Input tensor of shape (B, D).
- /// * `f` - Function to apply on the input tensor.
- ///
- /// # Returns
- ///
- /// The result of the function application, with a stochastic skip connection applied.
- #[inline]
- #[must_use]
- pub fn with_skip(
- &self,
- x: Tensor,
- f: F,
- ) -> Tensor
- where
- F: FnOnce(Tensor) -> Tensor,
- {
- x.clone() + self.forward(f(x))
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- use burn::backend::NdArray;
- use burn::prelude::Tensor;
- use burn::tensor::Distribution;
-
- #[test]
- fn test_drop_path() {
- let device = Default::default();
- let drop_prob = 0.5;
- let scale_by_keep = true;
-
- let config = DropPathConfig {
- drop_prob,
- scale_by_keep,
- };
-
- let module = config.init();
-
- let input =
- Tensor::::random([2, 3, 4, 5], Distribution::Uniform(0.0, 1.0), &device);
- let output = module.forward(input.clone());
-
- assert_eq!(input.dims(), output.dims());
- }
-
- #[test]
- fn test_drop_path_wrapper() {
- let device = Default::default();
-
- let n = 3;
- let shape = [n, 2, 4];
-
- let x = Tensor::::random(shape, Distribution::Uniform(0.0, 1.0), &device);
-
- // No-op case: not training and drop_prob = 0.0
- let training = false;
- let drop_prob = 0.0;
- let scale_by_keep = false;
- let res = drop_path(x.clone(), drop_prob, training, scale_by_keep);
- assert_eq!(res.dims(), x.dims());
- }
-
- #[test]
- fn test_drop_path_sample() {
- let device = Default::default();
-
- let n = 3;
- let shape = [n, 2, 4];
-
- let x = Tensor::::random(shape, Distribution::Uniform(0.0, 1.0), &device);
-
- // No-op case: not training and drop_prob = 0.0
- let training = false;
- let drop_prob = 0.0;
- let scale_by_keep = false;
- let res = _drop_path_sample(
- x.clone(),
- drop_prob,
- training,
- scale_by_keep,
- |shape, keep_prob, device| {
- assert_eq!(shape, [3, 1, 1]);
- assert_eq!(keep_prob, 1.0);
- Tensor::::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
- },
- );
- res.to_data().assert_eq(&x.clone().to_data(), true);
-
- // No-op case: training, but drop_prob = 0.0
- let training = true;
- let drop_prob = 0.0;
- let scale_by_keep = false;
- let res = _drop_path_sample(
- x.clone(),
- drop_prob,
- training,
- scale_by_keep,
- |shape, keep_prob, device| {
- assert_eq!(shape, [3, 1, 1]);
- assert_eq!(keep_prob, 1.0);
- Tensor::::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
- },
- );
- res.to_data().assert_eq(&x.clone().to_data(), true);
-
- // Training, but no scaling
- let training = true;
- let drop_prob = 0.5;
- let scale_by_keep = false;
- let res = _drop_path_sample(
- x.clone(),
- drop_prob,
- training,
- scale_by_keep,
- |shape, keep_prob, device| {
- assert_eq!(shape, [3, 1, 1]);
- assert_eq!(keep_prob, 0.5);
- Tensor::::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
- },
- );
- res.to_data().assert_eq(
- &(x.clone() * Tensor::::from_data([[[1.0]], [[0.0]], [[1.0]]], &device))
- .to_data(),
- true,
- );
-
- // Training, with scaling
- let training = true;
- let drop_prob = 0.5;
- let keep_prob = 1.0 - drop_prob;
- let scale_by_keep = true;
- let res = _drop_path_sample(
- x.clone(),
- drop_prob,
- training,
- scale_by_keep,
- |shape, keep_prob, device| {
- assert_eq!(shape, [3, 1, 1]);
- assert_eq!(keep_prob, 0.5);
- Tensor::::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
- },
- );
- res.to_data().assert_eq(
- &(x.clone() * Tensor::::from_data([[[1.0]], [[0.0]], [[1.0]]], &device))
- .div_scalar(keep_prob)
- .to_data(),
- true,
- );
- }
-
- #[test]
- fn test_droppath_module() {
- let drop_prob = 0.2;
- let config = DropPathConfig::new().with_drop_prob(drop_prob);
-
- assert_eq!(config.drop_prob(), 0.2);
- assert_eq!(config.keep_prob(), 1.0 - drop_prob);
- assert!(config.scale_by_keep());
-
- let module = config.init();
- assert_eq!(module.drop_prob(), 0.2);
- assert_eq!(module.keep_prob(), 1.0 - drop_prob);
- assert!(module.scale_by_keep());
-
- let device = Default::default();
- let shape = [2, 3, 4];
- let x = Tensor::::random(shape, Distribution::Uniform(0.0, 1.0), &device);
-
- // TODO(crutcher): work out how to enable/disable training mode in tests.
- let output = module.forward(x.clone());
- assert_eq!(x.dims(), output.dims());
- }
-}
diff --git a/crates/bimm/src/layers/drop/mod.rs b/crates/bimm/src/layers/drop/mod.rs
deleted file mode 100644
index ac22168a..00000000
--- a/crates/bimm/src/layers/drop/mod.rs
+++ /dev/null
@@ -1,5 +0,0 @@
-//! # Stochastic Drop Layers
-pub mod drop_block;
-pub mod drop_path;
-pub mod rate_table;
-mod size_config;
diff --git a/crates/bimm/src/layers/drop/rate_table.rs b/crates/bimm/src/layers/drop/rate_table.rs
deleted file mode 100644
index 362e8623..00000000
--- a/crates/bimm/src/layers/drop/rate_table.rs
+++ /dev/null
@@ -1,203 +0,0 @@
-//! Common rate table for `DropPath` regularization.
-
-use crate::compat::ops::float_vec_linspace;
-use alloc::vec::Vec;
-
-/// Computes a progressive incremental path drop rate for stochastic depth.
-///
-/// Used by SWIN Transformer V2 to apply a gradual increase in drop path rate.
-///
-/// # Arguments
-///
-/// * `drop_path_rate`: The final drop path rate to be achieved.
-/// * `depth`: The total number of layers in the model.
-///
-/// # Returns
-///
-/// A vector of drop path rates, one for each layer, starting from 0.0 and ending at `drop_path_rate`.
-#[inline(always)]
-#[must_use]
-pub fn progressive_dpr(
- drop_path_rate: f64,
- depth: usize,
-) -> Vec {
- float_vec_linspace(
- 0.0,
- drop_path_rate,
- // Total number of layers
- depth,
- )
-}
-
-/// Represents a table of progressive drop path rates for each layer in a model.
-pub struct DropPathRateDepthTable {
- progressive_dpr: Vec,
- layer_depths: Vec,
-}
-
-impl DropPathRateDepthTable {
- /// Creates a new `DropPathRateDepthTable` with the specified drop path rate and layer depths.
- ///
- /// # Arguments
- ///
- /// * `drop_path_rate`: The final drop path rate to be achieved.
- /// * `layer_depths`: A slice of layer depths, where each element represents the depth of a specific layer.
- ///
- /// # Returns
- ///
- /// A new `DropPathRateDepthTable` instance containing the progressive drop path rates for each layer.
- #[must_use]
- pub fn new(
- drop_path_rate: f64,
- layer_depths: &[usize],
- ) -> Self {
- let layer_depths = layer_depths.to_vec();
- let progressive_dpr = progressive_dpr(drop_path_rate, layer_depths.iter().sum());
- Self {
- progressive_dpr,
- layer_depths,
- }
- }
-
- /// Returns the depth of each layer.
- pub fn layer_depths(&self) -> &[usize] {
- &self.layer_depths
- }
-
- /// Returns the total depth of all layers.
- pub fn total_depth(&self) -> usize {
- self.layer_depths.iter().sum()
- }
-
- /// Returns the total number of layers.
- #[must_use]
- pub fn num_layers(&self) -> usize {
- self.layer_depths.len()
- }
-
- /// Returns the drop path rates for a specific layer.
- ///
- /// This will be a vector of size equal to the depth of the specified layer;
- /// corresponding to the progressive drop path rates for that layer.
- ///
- /// # Arguments
- ///
- /// * `layer_i`: The index of the layer for which to retrieve the drop path rates.
- ///
- /// # Panics
- ///
- /// If the layer index is out of bounds, it will panic with a message indicating the issue.
- #[must_use]
- pub fn layer_dprs(
- &self,
- layer_i: usize,
- ) -> Vec {
- if layer_i >= self.num_layers() {
- panic!(
- "Layer index {} out of bounds for {} layers",
- layer_i,
- self.num_layers()
- );
- }
- let depths = &self.layer_depths;
- let progressive_dpr1 = &self.progressive_dpr;
- let start = depths[..layer_i].iter().sum::();
- let end = start + depths[layer_i];
-
- progressive_dpr1[start..end].to_vec()
- }
-
- /// Returns the `layer_dprs` for all layers as a vector of vectors.
- #[inline(always)]
- #[must_use]
- pub fn layer_rates(&self) -> Vec> {
- (0..self.num_layers()).map(|i| self.layer_dprs(i)).collect()
- }
-
- /// Returns the progressive drop path rates for layer with the given depths.
- ///
- /// This is a convenience function that creates a new `DropPathRateDepthTable` and returns the layer rates.
- ///
- /// # Arguments
- ///
- /// * `drop_path_rate`: The final drop path rate to be achieved.
- /// * `layer_depths`: A slice of layer depths, where each element represents the depth of a specific layer.
- ///
- /// # Returns
- ///
- /// A vector of vectors, where each inner vector contains the progressive drop path rates for a specific layer.
- #[must_use]
- pub fn dpr_layer_rates(
- drop_path_rate: f64,
- layer_depths: &[usize],
- ) -> Vec> {
- let dpr_table = DropPathRateDepthTable::new(drop_path_rate, layer_depths);
- dpr_table.layer_rates()
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
- use crate::layers::drop::rate_table::progressive_dpr;
- use crate::testing::assert_close_to_vec;
- use alloc::vec;
- use hamcrest::prelude::*;
-
- #[test]
- fn test_incremental_drop_rate() {
- let drop_path_rate = 0.1;
- let depth = 9;
- let rates = progressive_dpr(drop_path_rate, depth);
- assert_close_to_vec(
- &rates,
- &[0.0, 0.0125, 0.025, 0.0375, 0.05, 0.0625, 0.075, 0.0875, 0.1],
- 0.001,
- );
- }
-
- #[test]
- fn test_table() {
- let depths = vec![2, 3, 4];
- let dpr_table = DropPathRateDepthTable::new(0.1, &depths);
-
- assert_eq!(dpr_table.total_depth(), 9);
-
- assert_that!(
- &dpr_table.layer_depths().to_vec(),
- contains(depths.clone()).exactly()
- );
-
- assert_close_to_vec(&dpr_table.layer_dprs(0), &[0.0, 0.0125], 0.001);
-
- assert_close_to_vec(&dpr_table.layer_dprs(1), &[0.025, 0.0375, 0.05], 0.001);
-
- let rates = dpr_table.layer_rates();
-
- assert_eq!(rates.len(), 3);
- assert_close_to_vec(&rates[0], &[0.0, 0.0125], 0.001);
- assert_close_to_vec(&rates[1], &[0.025, 0.0375, 0.05], 0.001);
- assert_close_to_vec(&rates[2], &[0.0625, 0.075, 0.0875, 0.1], 0.001);
- }
-
- #[should_panic(expected = "Layer index 3 out of bounds for 3 layers")]
- #[test]
- fn test_layer_dprs_out_of_bounds() {
- let depths = vec![2, 3, 4];
- let dpr_table = DropPathRateDepthTable::new(0.1, &depths);
- // This should panic because there are only 3 layers (0, 1, 2)
- let _d = dpr_table.layer_dprs(3);
- }
-
- #[test]
- fn test_dpr_layer_rates() {
- let drop_path_rate = 0.1;
- let layer_depths = vec![2, 3, 4];
- let rates = DropPathRateDepthTable::dpr_layer_rates(drop_path_rate, &layer_depths);
-
- assert_eq!(rates.len(), 3);
- assert_close_to_vec(&rates[0], &[0.0, 0.0125], 0.001);
- assert_close_to_vec(&rates[1], &[0.025, 0.0375, 0.05], 0.001);
- assert_close_to_vec(&rates[2], &[0.0625, 0.075, 0.0875, 0.1], 0.001);
- }
-}
diff --git a/crates/bimm/src/layers/drop/size_config.rs b/crates/bimm/src/layers/drop/size_config.rs
deleted file mode 100644
index 5508d501..00000000
--- a/crates/bimm/src/layers/drop/size_config.rs
+++ /dev/null
@@ -1,78 +0,0 @@
-use serde::{Deserialize, Serialize};
-
-/// Helper option to describe the size of a wrapper, relative to a wrapped object.
-///
-/// TODO: point this at `burn::...::SizeConfig` in "0.19.0"
-#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
-pub enum SizeConfig {
- /// Use the size of the source dataset.
- #[default]
- Default,
-
- /// Use the size as a ratio of the source dataset size.
- ///
- /// Must be >= 0.
- Ratio(f64),
-
- /// Use a fixed size.
- Fixed(usize),
-}
-
-impl SizeConfig {
- /// Construct a source which will have the same size as the source dataset.
- pub fn source() -> Self {
- Self::Default
- }
-
- /// Resolve the effective size.
- ///
- /// ## Arguments
- ///
- /// - `source_size`: the size of the source dataset.
- ///
- /// ## Returns
- ///
- /// The resolved size of the wrapper dataset.
- pub fn resolve(
- self,
- source_size: usize,
- ) -> usize {
- match self {
- SizeConfig::Default => source_size,
- SizeConfig::Ratio(ratio) => {
- assert!(ratio >= 0.0, "Ratio must be positive: {ratio}");
- ((source_size as f64) * ratio) as usize
- }
- SizeConfig::Fixed(size) => size,
- }
- }
-}
-
-impl From for SizeConfig {
- fn from(size: usize) -> Self {
- Self::Fixed(size)
- }
-}
-
-impl From for SizeConfig {
- fn from(ratio: f64) -> Self {
- Self::Ratio(ratio)
- }
-}
-
-#[cfg(test)]
-mod tests {
- use super::*;
-
- #[test]
- fn test_size_config() {
- assert_eq!(SizeConfig::default(), SizeConfig::Default);
-
- assert_eq!(SizeConfig::from(42), SizeConfig::Fixed(42));
-
- assert_eq!(SizeConfig::from(1.5), SizeConfig::Ratio(1.5));
-
- assert_eq!(SizeConfig::source(), SizeConfig::Default);
- assert_eq!(SizeConfig::source().resolve(50), 50);
- }
-}
diff --git a/crates/bimm/src/layers/mod.rs b/crates/bimm/src/layers/mod.rs
deleted file mode 100644
index f686a040..00000000
--- a/crates/bimm/src/layers/mod.rs
+++ /dev/null
@@ -1,6 +0,0 @@
-//! Common low-level modules for adding layers and operations in Burn.
-
-pub mod blocks;
-pub mod drop;
-pub mod patching;
-pub mod pool;
diff --git a/crates/bimm/src/layers/patching/mod.rs b/crates/bimm/src/layers/patching/mod.rs
deleted file mode 100644
index 25a41c43..00000000
--- a/crates/bimm/src/layers/patching/mod.rs
+++ /dev/null
@@ -1,3 +0,0 @@
-//! Spatial and temporal patching modules and operations.
-
-pub mod patch_embed;
diff --git a/crates/bimm/src/layers/patching/patch_embed.rs b/crates/bimm/src/layers/patching/patch_embed.rs
deleted file mode 100644
index 4fa0846b..00000000
--- a/crates/bimm/src/layers/patching/patch_embed.rs
+++ /dev/null
@@ -1,343 +0,0 @@
-//! # Patch Embedding layers and operations.
-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(
- &self,
- device: &B::Device,
- ) -> PatchEmbed {
- 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 {
- /// 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,
-
- /// Patch normalization layer, if enabled.
- pub norm: Option>,
-}
-
-impl PatchEmbedMeta for PatchEmbed {
- 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 PatchEmbed {
- /// 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,
- ) -> Tensor {
- assert_shape_contract_periodically!(
- ["batch", "d_input", "height", "width"],
- &x.dims(),
- &[
- ("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.dims(),
- &[
- ("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.dims(),
- &[
- ("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.dims(),
- &[
- ("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::(&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::(&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::