refactor(tensor): rename to_data_convert to to_data_as and into_data_convert to into_data_as, add to_vec_as and into_vec_as extensions for improved TensorData conversion methods

crates/bunsen-firehose-image/src/colortype_support.rs

@@ -260,43 +260,43 @@ mod tests {
 
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::L8),
-            DynamicImage::from(source.clone().to_luma8())
+            DynamicImage::from(source.to_luma8())
         );
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::La8),
-            DynamicImage::from(source.clone().to_luma_alpha8())
+            DynamicImage::from(source.to_luma_alpha8())
         );
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::Rgb8),
-            DynamicImage::from(source.clone().to_rgb8())
+            DynamicImage::from(source.to_rgb8())
         );
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::Rgba8),
-            DynamicImage::from(source.clone().to_rgba8())
+            DynamicImage::from(source.to_rgba8())
         );
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::L16),
-            DynamicImage::from(source.clone().to_luma16())
+            DynamicImage::from(source.to_luma16())
         );
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::La16),
-            DynamicImage::from(source.clone().to_luma_alpha16())
+            DynamicImage::from(source.to_luma_alpha16())
         );
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::Rgb16),
-            DynamicImage::from(source.clone().to_rgb16())
+            DynamicImage::from(source.to_rgb16())
         );
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::Rgba16),
-            DynamicImage::from(source.clone().to_rgba16())
+            DynamicImage::from(source.to_rgba16())
         );
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::Rgb32F),
-            DynamicImage::from(source.clone().to_rgb32f())
+            DynamicImage::from(source.to_rgb32f())
         );
         assert_eq!(
             convert_to_colortype(source.clone(), ColorType::Rgba32F),
-            DynamicImage::from(source.clone().to_rgba32f())
+            DynamicImage::from(source.to_rgba32f())
         );
     }
 }

crates/bunsen/src/blocks/images/drop/drop_path.rs

@@ -282,7 +282,7 @@ mod tests {
                 Tensor::<B, 3>::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
             },
         );
-        res.to_data().assert_eq(&x.clone().to_data(), true);
+        res.to_data().assert_eq(&x.to_data(), true);
 
         // No-op case: training, but drop_prob = 0.0
         let training = true;
@@ -299,7 +299,7 @@ mod tests {
                 Tensor::<B, 3>::from_data([[[1.0]], [[0.0]], [[1.0]]], device)
             },
         );
-        res.to_data().assert_eq(&x.clone().to_data(), true);
+        res.to_data().assert_eq(&x.to_data(), true);
 
         // Training, but no scaling
         let training = true;

crates/bunsen/src/blocks/rnn/lstm/ext_lstm_state.rs

@@ -267,8 +267,8 @@ mod tests {
         let device = Default::default();
         let state: ExtLstmState<B, 3> = random_state([2, 3, 4], &device);
 
-        let expected_cell = state.cell.clone().to_data();
-        let expected_hidden = state.hidden.clone().to_data();
+        let expected_cell = state.cell.to_data();
+        let expected_hidden = state.hidden.to_data();
 
         let (cell, hidden) = state.unpack();
 
@@ -354,7 +354,7 @@ mod tests {
         let device = Default::default();
         let state: ExtLstmState<B, 2> = random_state([2, 3], &device);
 
-        let expected_cell = state.cell.clone().to_data();
+        let expected_cell = state.cell.to_data();
 
         let roundtrip: ExtLstmState<B, 2> = state.unsqueeze_dim::<3>(1).squeeze_dim(1);
 
@@ -384,8 +384,8 @@ mod tests {
         let device = Default::default();
         let state: ExtLstmState<B, 2> = random_state([2, 3], &device);
 
-        let expected_cell = state.cell.clone().to_data();
-        let expected_hidden = state.hidden.clone().to_data();
+        let expected_cell = state.cell.to_data();
+        let expected_hidden = state.hidden.to_data();
 
         // The shape passed here is intentionally different; it must be ignored.
         let unwrapped = Some(state).unwrap_or_initial([9, 9], &device);

crates/bunsen/src/burner/tensor/tensor_op_ext.rs

@@ -12,6 +12,7 @@ use burn::{
         BasicOps,
         Bool,
         DType,
+        DataError,
         Element,
         Float,
         Int,
@@ -80,13 +81,19 @@ where
     K: BasicOps<B>,
     K::Elem: Element,
 {
-    /// Copies the current `Tensor` into `TensorData`; converts the dtype.
+    /// Copies the current `Tensor` into a `TensorData`; converts the dtype.
+    ///
+    /// By contract, this will yield the same result as
+    /// `tensor.to_data().convert::<E>(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;
+    fn to_data_as<E: Element>(&self) -> TensorData;
 
     /// Copies the current `Tensor` into `TensorData`; converts the dtype.
     ///
+    /// By contract, this will yield the same result as
+    /// `tensor.to_data().convert_dtype(dtype)`.
+    ///
     /// The conversion is a no-op if the dtype is the same as the current dtype.
     fn to_data_cast(
         &self,
@@ -95,11 +102,17 @@ where
 
     /// Converts the current `Tensor` into `TensorData`; converts the dtype.
     ///
+    /// By contract, this will yield the same result as
+    /// `tensor.into_data().convert::<E>(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;
+    fn into_data_as<E: Element>(self) -> TensorData;
 
     /// Converts the current `Tensor` into `TensorData`; converts the dtype.
     ///
+    /// By contract, this will yield the same result as
+    /// `tensor.into_data().convert_dtype(dtype)`.
+    ///
     /// The conversion is a no-op if the dtype is the same as the current dtype.
     fn into_data_cast(
         self,
@@ -113,8 +126,8 @@ where
     K: BasicOps<B>,
     K::Elem: Element,
 {
-    fn to_data_convert<E: Element>(&self) -> TensorData {
-        self.to_data().convert::<E>()
+    fn to_data_as<E: Element>(&self) -> TensorData {
+        self.to_data_cast(E::dtype())
     }
 
     fn to_data_cast(
@@ -124,8 +137,8 @@ where
         self.to_data().convert_dtype(dtype)
     }
 
-    fn into_data_convert<E: Element>(self) -> TensorData {
-        self.into_data().convert::<E>()
+    fn into_data_as<E: Element>(self) -> TensorData {
+        self.into_data_cast(E::dtype())
     }
 
     fn into_data_cast(
@@ -278,6 +291,74 @@ where
     }
 }
 
+/// Extension trait for `TensorData` that provides additional methods.
+pub trait TensorDataToVecAsExt {
+    /// Cast the data to a new dtype.
+    ///
+    /// TODO: Implement proper error handling in `TensorData`.
+    ///
+    /// # Returns
+    /// Ok(data) on success, (Currently) panics on failure.
+    fn try_cast(
+        self,
+        dtype: DType,
+    ) -> Result<TensorData, DataError>;
+
+    /// Convert the data to a new dtype.
+    ///
+    /// TODO: Implement proper error handling in `TensorData`.
+    ///
+    /// By contract, this is equivalent to:
+    /// `data.try_cast(E::dtype())`
+    ///
+    /// # Returns
+    /// Ok(data) on success, (Currently) panics on failure.
+    fn try_convert<E: Element>(self) -> Result<TensorData, DataError>;
+
+    /// Copy and convert the data to a [`Vec<E>`].
+    ///
+    /// By contract, this is equivalent to:
+    /// `data.clone().into_vec_as::<E>()`
+    ///
+    /// Particular conversions may provide more efficient implementations.
+    ///
+    /// # Returns
+    /// `Ok(vec)` on success, or an error if the conversion fails.
+    fn to_vec_as<E: Element>(&self) -> Result<Vec<E>, DataError>;
+
+    /// Convert the data to [`Vec<E>`].
+    ///
+    /// By contract, this is equivalent to:
+    /// `data.try_convert::<E>()?.to_vec::<E>()`
+    ///
+    /// Particular conversions may provide more efficient implementations.
+    ///
+    /// # Returns
+    /// `Ok(vec)` on success, or an error if the conversion fails.
+    fn into_vec_as<E: Element>(self) -> Result<Vec<E>, DataError>;
+}
+
+impl TensorDataToVecAsExt for TensorData {
+    fn try_cast(
+        self,
+        dtype: DType,
+    ) -> Result<TensorData, DataError> {
+        Ok(self.convert_dtype(dtype))
+    }
+
+    fn try_convert<E: Element>(self) -> Result<TensorData, DataError> {
+        self.try_cast(E::dtype())
+    }
+
+    fn to_vec_as<E: Element>(&self) -> Result<Vec<E>, DataError> {
+        self.clone().into_vec_as::<E>()
+    }
+
+    fn into_vec_as<E: Element>(self) -> Result<Vec<E>, DataError> {
+        self.try_convert::<E>()?.to_vec::<E>()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use burn::tensor::{
@@ -381,4 +462,41 @@ mod tests {
             .to_data()
             .assert_eq(&TensorData::from([3]), false);
     }
+
+    #[test]
+    fn test_to_vec_as() {
+        let data = TensorData::from([0.0f32, 1.0, 2.5]);
+
+        // Same-dtype copy.
+        assert_eq!(data.to_vec_as::<f32>().unwrap(), vec![0.0f32, 1.0, 2.5]);
+
+        // Widening cast (different element size).
+        assert_eq!(data.to_vec_as::<f64>().unwrap(), vec![0.0f64, 1.0, 2.5]);
+
+        // Float to int cast (same element size) truncates.
+        assert_eq!(data.to_vec_as::<i32>().unwrap(), vec![0i32, 1, 2]);
+
+        // The source data is borrowed, not consumed.
+        data.assert_eq(&TensorData::from([0.0f32, 1.0, 2.5]), true);
+    }
+
+    #[test]
+    fn test_into_vec_as() {
+        let data = TensorData::from([0i32, 1, 2, 3]);
+
+        // Same-dtype conversion.
+        assert_eq!(
+            data.clone().into_vec_as::<i32>().unwrap(),
+            vec![0i32, 1, 2, 3]
+        );
+
+        // Int to float cast.
+        assert_eq!(
+            data.clone().into_vec_as::<f32>().unwrap(),
+            vec![0.0f32, 1.0, 2.0, 3.0]
+        );
+
+        // Narrowing int cast.
+        assert_eq!(data.into_vec_as::<u8>().unwrap(), vec![0u8, 1, 2, 3]);
+    }
 }

crates/bunsen/src/kits/bimm/resnet/blocks/layer_block.rs

@@ -467,6 +467,7 @@ mod tests {
     use crate::{
         contracts::assert_shape_contract,
         kits::bimm::resnet::blocks::BasicBlockConfig,
+        prelude::*,
         support::testing::PerformanceBackend,
     };
 
@@ -556,7 +557,7 @@ mod tests {
             expected = block.forward(expected);
         }
         output
-            .to_data()
-            .assert_approx_eq::<F>(&expected.to_data(), Tolerance::default());
+            .to_data_as::<F>()
+            .assert_approx_eq::<F>(&expected.to_data_as::<F>(), Tolerance::default());
     }
 }

crates/bunsen/src/kits/bimm/swin/v2/blocks/window_attention/pos_grid.rs

@@ -284,7 +284,7 @@ mod tests {
 
         let device = Default::default();
         let rel = window_attention_relative_position_index::<B>(window_shape, &device);
-        rel.clone().to_data().assert_eq(
+        rel.to_data().assert_eq(
             &TensorData::from([
                 [7, 6, 5, 2, 1, 0],
                 [8, 7, 6, 3, 2, 1],

crates/bunsen/src/kits/sims/conway/life2d.rs

@@ -88,7 +88,7 @@ where
 
     state
         .slice(slices)
-        .to_data_convert::<bool>()
+        .to_data_as::<bool>()
         .to_vec()
         .unwrap()
         .chunks(w)

crates/bunsen/src/kits/sims/lbm/d2q9/collision.rs

@@ -119,7 +119,7 @@ mod tests {
         // Invariant: density(collision(dist, param)) == density(dist)
         density(col_dist.clone())
             .to_data()
-            .assert_approx_eq::<f32>(&rho.clone().to_data(), Tolerance::default());
+            .assert_approx_eq::<f32>(&rho.to_data(), Tolerance::default());
 
         // With Correction
         {

crates/bunsen/src/kits/sims/lbm/d2q9/space.rs

@@ -395,7 +395,7 @@ mod tests {
 
         let momentum = macroscopic_momentum(dist.clone(), lbm_tables.e_vec());
 
-        momentum.clone().to_data().assert_approx_eq::<f32>(
+        momentum.to_data().assert_approx_eq::<f32>(
             &Tensor::<B, 3>::from_data([[[-1., -1.], [15., 5.]]], &device).to_data(),
             Tolerance::default(),
         );
@@ -404,7 +404,7 @@ mod tests {
 
         let rho_data = rho.to_data().to_vec::<f32>().unwrap();
 
-        u.clone().to_data().assert_approx_eq::<f32>(
+        u.to_data().assert_approx_eq::<f32>(
             &Tensor::<B, 3>::from_data(
                 [[
                     [-1. / rho_data[0], -1. / rho_data[0]],
@@ -418,7 +418,7 @@ mod tests {
 
         let v_sq = velocity_squared(u.clone());
 
-        v_sq.clone().to_data().assert_approx_eq::<f32>(
+        v_sq.to_data().assert_approx_eq::<f32>(
             &Tensor::<B, 2>::from_data(
                 [[
                     (1. + 1.) / rho_data[0].powi(2),

crates/bunsen/src/kits/sims/lbm/d2q9/thermal.rs

@@ -165,7 +165,7 @@ mod tests {
 
         let parts = ldv_projection(e.clone(), u.clone());
 
-        parts.clone().to_data().assert_approx_eq::<f32>(
+        parts.to_data().assert_approx_eq::<f32>(
             &Tensor::<B, 5>::from_data(
                 [[
                     [
@@ -187,7 +187,7 @@ mod tests {
 
         let e_u = lattice_dot_velocity(u.clone(), e.clone());
 
-        e_u.clone().to_data().assert_approx_eq::<f32>(
+        e_u.to_data().assert_approx_eq::<f32>(
             &parts.sum_dim(4).squeeze_dims::<4>(&[4]).to_data(),
             Tolerance::default(),
         );