refactor(tensor): unify to_data_convert and to_data_cast methods across tensor operations, update usage in lbm2d_vis, conway_vis, and life2d implementations for consistency

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<B, const D: usize, K>
+where
+    B: Backend,
+    K: BasicOps<B>,
+    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<E: Element>(&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<E: Element>(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<B, const D: usize, K> TensorElemOpExt<B, D, K> for Tensor<B, D, K>
+where
+    B: Backend,
+    K: BasicOps<B>,
+    K::Elem: Element,
+{
+    fn to_data_convert<E: Element>(&self) -> TensorData {
+        self.to_data().convert::<E>()
+    }
+
+    fn to_data_cast(
+        &self,
+        dtype: DType,
+    ) -> TensorData {
+        self.to_data().convert_dtype(dtype)
+    }
+
+    fn into_data_convert<E: Element>(self) -> TensorData {
+        self.into_data().convert::<E>()
+    }
+
+    fn into_data_cast(
+        self,
+        dtype: DType,
+    ) -> TensorData {
+        self.into_data().convert_dtype(dtype)
+    }
+}
+
 /// Tensor Extension trait for ordered operations.
 pub trait TensorOrderedOpExt<B, const D: usize, K>
 where

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::<i32>().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::<Vec<_>>(),
-        );
-    }
-
-    result
+    let [_, w] = slices_shape(&slices);
+
+    state
+        .slice(slices)
+        .to_data_convert::<bool>()
+        .to_vec()
+        .unwrap()
+        .chunks(w)
+        .map(<[_]>::to_vec)
+        .collect()
 }
 
 /// Returns the next board.

crates/bunsen/src/kits/speech/silero_vad/blocks/module.rs

@@ -477,7 +477,10 @@ impl<B: Backend> SileroVad<B> {
 
         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<B: Backend> SileroVad<B> {
 
         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())]

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::<bool>();
+                    let frame = conway.state.clone().into_data_convert::<bool>();
                     *frame_handle_1.lock().unwrap() = frame;
 
                     t1 = std::time::Instant::now();

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<B: Backend>(
             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::<f32>();
+            *vis_cells_publish.lock().unwrap() = cells.to_data_convert::<f32>();
 
             last_export = std::time::Instant::now();
         }
@@ -207,7 +211,7 @@ fn run<B: Backend>(
     let mut app = FlowVisApp {
         gl: GlGraphics::new(opengl),
         cell_data: vis_cells,
-        solid_mask: solid_mask.to_data().convert::<bool>(),
+        solid_mask: solid_mask.to_data_convert::<bool>(),
         opacity: args.opacity,
     };
 

examples/resnet_tiny/src/main.rs

@@ -199,7 +199,7 @@ fn main() -> anyhow::Result<()> {
             type B = burn::backend::Wgpu<burn::tensor::bf16>;
         }
         _ => {
-            type B =burn::backend::Flex;
+            type B = burn::backend::Flex;
         }
     }
     backend_main::<Autodiff<B>>(&args)

examples/whisper-dev/src/main.rs

@@ -21,18 +21,10 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     println!("{:#?}", args);
 
     cfg_select! {
-        feature = "cuda" => {
-            run::<burn::backend::cuda::Cuda>(args)
-        }
-        feature = "metal" => {
-            run::<burn::backend::Metal>(args)
-        }
-        feature = "wgpu" => {
-            run::<burn::backend::wgpu::Wgpu>(args)
-        }
-        feature = "flex" => {
-            run::<burn::backend::flex::Flex>(args)
-        }
+        feature = "cuda" => run::<burn::backend::cuda::Cuda>(args),
+        feature = "metal" => run::<burn::backend::Metal>(args),
+        feature = "wgpu" => run::<burn::backend::wgpu::Wgpu>(args),
+        feature = "flex" => run::<burn::backend::flex::Flex>(args),
         _ => {
             compile_error!("No Backend enabled");
         }