diff --git a/crates/bunsen/src/burner/tensor/tensor_op_ext.rs b/crates/bunsen/src/burner/tensor/tensor_op_ext.rs index 5f68ac5..e6ab0b2 100644 --- a/crates/bunsen/src/burner/tensor/tensor_op_ext.rs +++ b/crates/bunsen/src/burner/tensor/tensor_op_ext.rs @@ -5,11 +5,14 @@ use burn::{ prelude::{ Backend, ElementConversion, + TensorData, }, tensor::{ AsIndex, BasicOps, Bool, + DType, + Element, Float, Int, }, @@ -70,6 +73,69 @@ where } } +/// [`Tensor`] element-type aware extension operations. +pub trait TensorElemOpExt +where + B: Backend, + K: BasicOps, + K::Elem: Element, +{ + /// Copies the current `Tensor` into `TensorData`; converts the dtype. + /// + /// The conversion is a no-op if the dtype is the same as the current dtype. + fn to_data_convert(&self) -> TensorData; + + /// Copies the current `Tensor` into `TensorData`; converts the dtype. + /// + /// The conversion is a no-op if the dtype is the same as the current dtype. + fn to_data_cast( + &self, + dtype: DType, + ) -> TensorData; + + /// Converts the current `Tensor` into `TensorData`; converts the dtype. + /// + /// The conversion is a no-op if the dtype is the same as the current dtype. + fn into_data_convert(self) -> TensorData; + + /// Converts the current `Tensor` into `TensorData`; converts the dtype. + /// + /// The conversion is a no-op if the dtype is the same as the current dtype. + fn into_data_cast( + self, + dtype: DType, + ) -> TensorData; +} + +impl TensorElemOpExt for Tensor +where + B: Backend, + K: BasicOps, + K::Elem: Element, +{ + fn to_data_convert(&self) -> TensorData { + self.to_data().convert::() + } + + fn to_data_cast( + &self, + dtype: DType, + ) -> TensorData { + self.to_data().convert_dtype(dtype) + } + + fn into_data_convert(self) -> TensorData { + self.into_data().convert::() + } + + fn into_data_cast( + self, + dtype: DType, + ) -> TensorData { + self.into_data().convert_dtype(dtype) + } +} + /// Tensor Extension trait for ordered operations. pub trait TensorOrderedOpExt where diff --git a/crates/bunsen/src/kits/sims/conway/life2d.rs b/crates/bunsen/src/kits/sims/conway/life2d.rs index 1311861..38ca63d 100644 --- a/crates/bunsen/src/kits/sims/conway/life2d.rs +++ b/crates/bunsen/src/kits/sims/conway/life2d.rs @@ -11,13 +11,15 @@ use burn::{ s, }, tensor::{ - DType::I32, Distribution, Slice, }, }; -use crate::prelude::TensorBoolOpExt; +use crate::prelude::{ + TensorBoolOpExt, + TensorElemOpExt, +}; /// Fuzzes the state. /// @@ -82,23 +84,16 @@ where R: SliceArg, { let slices: [Slice; 2] = ranges.into_slices(&state.shape()).try_into().unwrap(); - let [h, w] = slices_shape(&slices); - - let block_data = state.slice(slices).int().cast(I32).to_data(); - let block_data = block_data.to_vec::().unwrap(); - - let mut result = Vec::with_capacity(h); - for hidx in 0..h { - let start = hidx * w; - result.push( - block_data[start..start + w] - .iter() - .map(|&b| b != 0) - .collect::>(), - ); - } - - result + let [_, w] = slices_shape(&slices); + + state + .slice(slices) + .to_data_convert::() + .to_vec() + .unwrap() + .chunks(w) + .map(<[_]>::to_vec) + .collect() } /// Returns the next board. diff --git a/crates/bunsen/src/kits/speech/silero_vad/blocks/module.rs b/crates/bunsen/src/kits/speech/silero_vad/blocks/module.rs index fa699ba..df68689 100644 --- a/crates/bunsen/src/kits/speech/silero_vad/blocks/module.rs +++ b/crates/bunsen/src/kits/speech/silero_vad/blocks/module.rs @@ -477,7 +477,10 @@ impl SileroVad { cfg_select! { any(test, debug_assertions) => { - use crate::contracts::{unpack_shape_contract, assert_shape_contract_periodically}; + use crate::contracts::{ + assert_shape_contract_periodically, + unpack_shape_contract, + }; let [steps, batch] = unpack_shape_contract!( ["steps", "batch", "samples"], &chunk_seq, @@ -558,9 +561,12 @@ impl SileroVad { cfg_select! { any(test, debug_assertions) => { - use crate::contracts::{unpack_shape_contract, assert_shape_contract_periodically}; + use crate::contracts::{ + assert_shape_contract_periodically, + unpack_shape_contract, + }; let [batch] = unpack_shape_contract!( - [ "batch", "samples"], + ["batch", "samples"], &chunk, &["batch"], &[("samples", self.chunk_size())] diff --git a/examples/conway_vis/src/main.rs b/examples/conway_vis/src/main.rs index c7493a8..77cdcee 100644 --- a/examples/conway_vis/src/main.rs +++ b/examples/conway_vis/src/main.rs @@ -18,6 +18,10 @@ use bunsen::{ ConwayLife2DConfig, ConwayLife2DState, }, + prelude::{ + TensorElemOpExt, + TensorOpExt, + }, support::validators::parse_grid_shape, zspace::ravel_dims, }; @@ -280,7 +284,7 @@ impl Simulation { if t1 - last_export > export_duration { last_export = t1; - let frame = conway.state.clone().into_data().convert::(); + let frame = conway.state.clone().into_data_convert::(); *frame_handle_1.lock().unwrap() = frame; t1 = std::time::Instant::now(); diff --git a/examples/lbm2d_vis/src/main.rs b/examples/lbm2d_vis/src/main.rs index e7f1a12..a5e8753 100644 --- a/examples/lbm2d_vis/src/main.rs +++ b/examples/lbm2d_vis/src/main.rs @@ -24,6 +24,10 @@ use bunsen::{ SPEED_OF_SOUND, macroscopic_momentum, }, + prelude::{ + TensorElemOpExt, + TensorOpExt, + }, support::validators::parse_grid_shape, }; use burn::{ @@ -185,7 +189,7 @@ fn run( let cells = ((cells / scale) + 1.0) / 2.0; // let cells = cells.mul_scalar(std::f64::consts::PI / 2.0).sin(); - *vis_cells_publish.lock().unwrap() = cells.cast(DType::F32).to_data().convert::(); + *vis_cells_publish.lock().unwrap() = cells.to_data_convert::(); last_export = std::time::Instant::now(); } @@ -207,7 +211,7 @@ fn run( let mut app = FlowVisApp { gl: GlGraphics::new(opengl), cell_data: vis_cells, - solid_mask: solid_mask.to_data().convert::(), + solid_mask: solid_mask.to_data_convert::(), opacity: args.opacity, }; diff --git a/examples/resnet_tiny/src/main.rs b/examples/resnet_tiny/src/main.rs index 4c0d713..7293a0d 100644 --- a/examples/resnet_tiny/src/main.rs +++ b/examples/resnet_tiny/src/main.rs @@ -199,7 +199,7 @@ fn main() -> anyhow::Result<()> { type B = burn::backend::Wgpu; } _ => { - type B =burn::backend::Flex; + type B = burn::backend::Flex; } } backend_main::>(&args) diff --git a/examples/whisper-dev/src/main.rs b/examples/whisper-dev/src/main.rs index 0b4bc8a..1013abd 100644 --- a/examples/whisper-dev/src/main.rs +++ b/examples/whisper-dev/src/main.rs @@ -21,18 +21,10 @@ fn main() -> Result<(), Box> { println!("{:#?}", args); cfg_select! { - feature = "cuda" => { - run::(args) - } - feature = "metal" => { - run::(args) - } - feature = "wgpu" => { - run::(args) - } - feature = "flex" => { - run::(args) - } + feature = "cuda" => run::(args), + feature = "metal" => run::(args), + feature = "wgpu" => run::(args), + feature = "flex" => run::(args), _ => { compile_error!("No Backend enabled"); }