torch

crutcher · crutcher · commit f51994083d50 · 2025-09-23T09:45:11.000-07:00
diff --git a/crates/burn-autodiff/src/ops/bool_tensor.rs b/crates/burn-autodiff/src/ops/bool_tensor.rs
@@ -108,7 +108,12 @@ impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {
         B::bool_repeat_dim(tensor, dim, times)
     }
 
-    fn bool_unfold(tensor: BoolTensor<Self>, dim: usize, size: usize, step: usize) -> BoolTensor<Self> {
+    fn bool_unfold(
+        tensor: BoolTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> BoolTensor<Self> {
         B::bool_unfold(tensor, dim, size, step)
     }
 }
diff --git a/crates/burn-autodiff/src/ops/int_tensor.rs b/crates/burn-autodiff/src/ops/int_tensor.rs
@@ -378,7 +378,12 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_cast(tensor, dtype)
     }
 
-    fn int_unfold(tensor: IntTensor<Self>, dim: usize, size: usize, step: usize) -> IntTensor<Self> {
+    fn int_unfold(
+        tensor: IntTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> IntTensor<Self> {
         B::int_unfold(tensor, dim, size, step)
     }
 }
diff --git a/crates/burn-autodiff/src/ops/tensor.rs b/crates/burn-autodiff/src/ops/tensor.rs
@@ -2593,7 +2593,12 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
     // TODO: Implement float_prod and float_sum
     // https://github.com/tracel-ai/burn/issues/1458
 
-    fn float_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {
+    fn float_unfold(
+        tensor: FloatTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> FloatTensor<Self> {
         AutodiffTensor::new(B::float_unfold(tensor.primitive, dim, size, step))
     }
 }
diff --git a/crates/burn-cubecl/src/ops/base.rs b/crates/burn-cubecl/src/ops/base.rs
@@ -1,8 +1,11 @@
-use std::cmp::max;
 use crate::{CubeRuntime, element::CubeElement, kernel, tensor::CubeTensor};
 use burn_common::tensor::{ReshapeAction, reshape_action};
-use burn_tensor::{Shape, TensorData, quantization::{QTensorPrimitive, QuantLevel}};
+use burn_tensor::{
+    Shape, TensorData,
+    quantization::{QTensorPrimitive, QuantLevel},
+};
 use cubecl::{server::CopyDescriptor, tensor_vectorization_factor};
+use std::cmp::max;
 
 pub(crate) fn from_data<R: CubeRuntime>(data: TensorData, device: &R::Device) -> CubeTensor<R> {
     let shape: Shape = (&data.shape).into();
diff --git a/crates/burn-cubecl/src/ops/bool_ops.rs b/crates/burn-cubecl/src/ops/bool_ops.rs
@@ -142,7 +142,12 @@ where
         kernel::flip::<R, BT, BT>(tensor, axes)
     }
 
-    fn bool_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {
+    fn bool_unfold(
+        tensor: FloatTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> FloatTensor<Self> {
         unfold(tensor, dim, size, step)
     }
 }
diff --git a/crates/burn-cubecl/src/ops/float_ops.rs b/crates/burn-cubecl/src/ops/float_ops.rs
@@ -703,7 +703,12 @@ where
         }
     }
 
-    fn float_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {
+    fn float_unfold(
+        tensor: FloatTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> FloatTensor<Self> {
         unfold(tensor, dim, size, step)
     }
 }
diff --git a/crates/burn-cubecl/src/ops/int_ops.rs b/crates/burn-cubecl/src/ops/int_ops.rs
@@ -685,7 +685,12 @@ where
         )
     }
 
-    fn int_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {
+    fn int_unfold(
+        tensor: FloatTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> FloatTensor<Self> {
         unfold(tensor, dim, size, step)
     }
 }
diff --git a/crates/burn-fusion/src/ops/boolean.rs b/crates/burn-fusion/src/ops/boolean.rs
@@ -1,9 +1,13 @@
-use std::cmp::max;
-use burn_ir::{BaseOperationIr, BinaryOpIr, BoolOperationIr, CatOpIr, ExpandOpIr, FlipOpIr, HandleContainer, InitOperationIr, OperationIr, PermuteOpIr, RepeatDimOpIr, SliceAssignOpIr, SliceOpIr, SwapDimsOpIr, TensorIr, UnaryOpIr, UnfoldOpIr};
+use burn_ir::{
+    BaseOperationIr, BinaryOpIr, BoolOperationIr, CatOpIr, ExpandOpIr, FlipOpIr, HandleContainer,
+    InitOperationIr, OperationIr, PermuteOpIr, RepeatDimOpIr, SliceAssignOpIr, SliceOpIr,
+    SwapDimsOpIr, TensorIr, UnaryOpIr, UnfoldOpIr,
+};
 use burn_tensor::{
     Device, Element, Shape, Slice, TensorData, TensorMetadata,
     ops::{BoolTensor, BoolTensorOps, FloatTensor, IntTensor, binary_ops_shape},
 };
+use std::cmp::max;
 use std::marker::PhantomData;
 
 use crate::{
@@ -742,7 +746,12 @@ impl<B: FusionBackend> BoolTensorOps<Self> for Fusion<B> {
         out
     }
 
-    fn bool_unfold(tensor: BoolTensor<Self>, dim: usize, size: usize, step: usize) -> BoolTensor<Self> {
+    fn bool_unfold(
+        tensor: BoolTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> BoolTensor<Self> {
         #[derive(new, Debug)]
         struct UnfoldOps<B: FusionBackend> {
             desc: UnfoldOpIr,
@@ -752,11 +761,7 @@ impl<B: FusionBackend> BoolTensorOps<Self> for Fusion<B> {
         impl<B: FusionBackend> Operation<B::FusionRuntime> for UnfoldOps<B> {
             fn execute(&self, handles: &mut HandleContainer<B::Handle>) {
                 let input = handles.get_bool_tensor::<B>(&self.desc.input);
-                let output = B::bool_unfold(
-                    input,
-                    self.desc.dim,
-                    self.desc.size,
-                    self.desc.step);
+                let output = B::bool_unfold(input, self.desc.dim, self.desc.size, self.desc.step);
 
                 handles.register_bool_tensor::<B>(&self.desc.out.id, output);
             }
diff --git a/crates/burn-fusion/src/ops/float.rs b/crates/burn-fusion/src/ops/float.rs
@@ -1,3 +1,4 @@
+use super::NoOp;
 use crate::{
     Fusion, FusionBackend, binary_float_cmp_ops, binary_float_ops,
     client::FusionClient,
@@ -12,9 +13,8 @@ use burn_tensor::{
     Device, Distribution, Element, FloatDType, Shape, Slice, TensorData, TensorMetadata,
     ops::{BoolTensor, FloatElem, FloatTensor, FloatTensorOps, IntTensor, binary_ops_shape},
 };
-use std::{marker::PhantomData, ops::Range};
 use std::cmp::max;
-use super::NoOp;
+use std::{marker::PhantomData, ops::Range};
 
 impl<B: FusionBackend> FloatTensorOps<Self> for Fusion<B> {
     fn float_from_data(data: TensorData, device: &Device<Self>) -> FloatTensor<Self> {
@@ -2260,7 +2260,12 @@ impl<B: FusionBackend> FloatTensorOps<Self> for Fusion<B> {
         out
     }
 
-    fn float_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {
+    fn float_unfold(
+        tensor: FloatTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> FloatTensor<Self> {
         #[derive(new, Debug)]
         struct UnfoldOps<B: FusionBackend> {
             desc: UnfoldOpIr,
@@ -2270,11 +2275,7 @@ impl<B: FusionBackend> FloatTensorOps<Self> for Fusion<B> {
         impl<B: FusionBackend> Operation<B::FusionRuntime> for UnfoldOps<B> {
             fn execute(&self, handles: &mut HandleContainer<B::Handle>) {
                 let input = handles.get_float_tensor::<B>(&self.desc.input);
-                let output = B::float_unfold(
-                    input,
-                    self.desc.dim,
-                    self.desc.size,
-                    self.desc.step);
+                let output = B::float_unfold(input, self.desc.dim, self.desc.size, self.desc.step);
 
                 handles.register_float_tensor::<B>(&self.desc.out.id, output);
             }
diff --git a/crates/burn-fusion/src/ops/int.rs b/crates/burn-fusion/src/ops/int.rs
@@ -2173,7 +2173,12 @@ impl<B: FusionBackend> IntTensorOps<Self> for Fusion<B> {
         out
     }
 
-    fn int_unfold(tensor: IntTensor<Self>, dim: usize, size: usize, step: usize) -> IntTensor<Self> {
+    fn int_unfold(
+        tensor: IntTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> IntTensor<Self> {
         #[derive(new, Debug)]
         struct UnfoldOps<B: FusionBackend> {
             desc: UnfoldOpIr,
@@ -2183,11 +2188,7 @@ impl<B: FusionBackend> IntTensorOps<Self> for Fusion<B> {
         impl<B: FusionBackend> Operation<B::FusionRuntime> for UnfoldOps<B> {
             fn execute(&self, handles: &mut HandleContainer<B::Handle>) {
                 let input = handles.get_int_tensor::<B>(&self.desc.input);
-                let output = B::int_unfold(
-                    input,
-                    self.desc.dim,
-                    self.desc.size,
-                    self.desc.step);
+                let output = B::int_unfold(input, self.desc.dim, self.desc.size, self.desc.step);
 
                 handles.register_int_tensor::<B>(&self.desc.out.id, output);
             }
diff --git a/crates/burn-ndarray/src/ops/base.rs b/crates/burn-ndarray/src/ops/base.rs
@@ -221,10 +221,14 @@ where
     ///
     /// A tensor view with shape ``[pre=..., windows, size, post=...]``.
     #[allow(unused)]
-    pub(crate) fn unfold(tensor: SharedArray<E>, dim: usize, size: usize, step: usize) -> SharedArray<E> {
+    pub(crate) fn unfold(
+        tensor: SharedArray<E>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> SharedArray<E> {
         unimplemented!("ndarray unfold not implemented yet");
     }
-
 }
 
 #[cfg(feature = "simd")]
diff --git a/crates/burn-ndarray/src/ops/bool_tensor.rs b/crates/burn-ndarray/src/ops/bool_tensor.rs
@@ -156,7 +156,12 @@ where
     }
 
     #[allow(unused)]
-    fn bool_unfold(tensor: BoolTensor<Self>, dim: usize, size: usize, step: usize) -> BoolTensor<Self> {
+    fn bool_unfold(
+        tensor: BoolTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> BoolTensor<Self> {
         todo!()
     }
 }
diff --git a/crates/burn-ndarray/src/ops/int_tensor.rs b/crates/burn-ndarray/src/ops/int_tensor.rs
@@ -447,7 +447,12 @@ where
     }
 
     #[allow(unused)]
-    fn int_unfold(tensor: IntTensor<Self>, dim: usize, size: usize, step: usize) -> IntTensor<Self> {
+    fn int_unfold(
+        tensor: IntTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> IntTensor<Self> {
         todo!()
     }
 }
diff --git a/crates/burn-ndarray/src/ops/tensor.rs b/crates/burn-ndarray/src/ops/tensor.rs
@@ -499,7 +499,12 @@ where
     }
 
     #[allow(unused)]
-    fn float_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {
+    fn float_unfold(
+        tensor: FloatTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> FloatTensor<Self> {
         execute_with_float_dtype!(tensor, |tensor| NdArrayOps::unfold(tensor, dim, size, step))
     }
 }
diff --git a/crates/burn-tch/src/ops/base.rs b/crates/burn-tch/src/ops/base.rs
@@ -447,6 +447,19 @@ impl TchOps {
         TchTensor::from_existing(broadcasted_tensor, storage)
     }
 
+    pub fn unfold(tensor: TchTensor, dim: usize, size: usize, step: usize) -> TchTensor {
+        let storage = tensor.storage.clone();
+        let uf_tensor = tensor.tensor.unfold(dim as i64, size as i64, step as i64);
+
+        let mut perm: Vec<i64> = (0..uf_tensor.dim()).into_iter().map(|x| x as i64).collect();
+        let last = perm.remove(perm.len() - 1);
+        perm.insert(dim + 1, last);
+
+        let uf_tensor = uf_tensor.permute(perm);
+
+        TchTensor::from_existing(uf_tensor, storage)
+    }
+
     pub fn sort(tensor: TchTensor, dim: usize, descending: bool) -> TchTensor {
         TchTensor::new(tensor.tensor.sort(dim as i64, descending).0)
     }
diff --git a/crates/burn-tch/src/ops/bool_tensor.rs b/crates/burn-tch/src/ops/bool_tensor.rs
@@ -1,5 +1,6 @@
 use super::TchOps;
 use crate::{LibTorch, LibTorchDevice, TchShape, TchTensor, element::TchElement};
+use burn_tensor::ops::IntTensor;
 use burn_tensor::{Shape, TensorData, TensorMetadata, backend::Backend, ops::BoolTensorOps};
 use std::ops::Range;
 
@@ -150,4 +151,13 @@ impl<E: TchElement> BoolTensorOps<Self> for LibTorch<E> {
     fn bool_expand(tensor: TchTensor, shape: Shape) -> TchTensor {
         TchOps::expand(tensor, shape)
     }
+
+    fn bool_unfold(
+        tensor: IntTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> IntTensor<Self> {
+        TchOps::unfold(tensor, dim, size, step)
+    }
 }
diff --git a/crates/burn-tch/src/ops/int_tensor.rs b/crates/burn-tch/src/ops/int_tensor.rs
@@ -499,4 +499,13 @@ impl<E: TchElement> IntTensorOps<Self> for LibTorch<E> {
             TchTensor::new(tensor.tensor.to_kind(kind))
         }
     }
+
+    fn int_unfold(
+        tensor: IntTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> IntTensor<Self> {
+        TchOps::unfold(tensor, dim, size, step)
+    }
 }
diff --git a/crates/burn-tch/src/ops/tensor.rs b/crates/burn-tch/src/ops/tensor.rs
@@ -1,5 +1,6 @@
 use super::TchOps;
 use crate::{IntoKind, LibTorch, LibTorchDevice, TchShape, TchTensor, element::TchElement};
+use burn_tensor::ops::FloatTensor;
 use burn_tensor::{
     DType, Distribution, ElementConversion, FloatDType, Shape, TensorData, TensorMetadata,
     backend::Backend,
@@ -473,4 +474,13 @@ impl<E: TchElement> FloatTensorOps<Self> for LibTorch<E> {
             TchTensor::new(tensor.tensor.to_kind(kind))
         }
     }
+
+    fn float_unfold(
+        tensor: FloatTensor<Self>,
+        dim: usize,
+        size: usize,
+        step: usize,
+    ) -> FloatTensor<Self> {
+        TchOps::unfold(tensor, dim, size, step)
+    }
 }
diff --git a/crates/burn-tensor/src/tensor/api/base.rs b/crates/burn-tensor/src/tensor/api/base.rs
@@ -2197,7 +2197,12 @@ where
     /// # Returns
     ///
     /// A tensor view with shape ``[pre=..., windows, size, post=...]``.
-    pub fn unfold<const D2: usize, I: AsIndex>(self, dim: I, size: usize, step: usize) -> Tensor<B, D2, K> {
+    pub fn unfold<const D2: usize, I: AsIndex>(
+        self,
+        dim: I,
+        size: usize,
+        step: usize,
+    ) -> Tensor<B, D2, K> {
         let dim = canonicalize_dim(dim, D, false);
         Tensor::<B, D2, K>::new(K::unfold(self.primitive, dim, size, step))
     }
diff --git a/crates/burn-tensor/src/tensor/ops/bool_tensor.rs b/crates/burn-tensor/src/tensor/ops/bool_tensor.rs
@@ -437,10 +437,5 @@ pub trait BoolTensorOps<B: Backend> {
     /// # Returns
     ///
     /// A tensor view with shape ``[pre=..., windows, size, post=...]``.
-    fn bool_unfold(
-        tensor: BoolTensor<B>,
-        dim: usize,
-        size: usize,
-        step: usize,
-    ) -> BoolTensor<B>;
+    fn bool_unfold(tensor: BoolTensor<B>, dim: usize, size: usize, step: usize) -> BoolTensor<B>;
 }
diff --git a/crates/burn-tensor/src/tensor/ops/int_tensor.rs b/crates/burn-tensor/src/tensor/ops/int_tensor.rs
@@ -1228,10 +1228,5 @@ pub trait IntTensorOps<B: Backend> {
     /// # Returns
     ///
     /// A tensor view with shape ``[pre=..., windows, size, post=...]``.
-    fn int_unfold(
-        tensor: IntTensor<B>,
-        dim: usize,
-        size: usize,
-        step: usize,
-    ) -> IntTensor<B>;
+    fn int_unfold(tensor: IntTensor<B>, dim: usize, size: usize, step: usize) -> IntTensor<B>;
 }
diff --git a/crates/burn-tensor/src/tensor/ops/tensor.rs b/crates/burn-tensor/src/tensor/ops/tensor.rs
@@ -1403,10 +1403,6 @@ pub trait FloatTensorOps<B: Backend> {
     /// # Returns
     ///
     /// A tensor view with shape ``[pre=..., windows, size, post=...]``.
-    fn float_unfold(
-        tensor: FloatTensor<B>,
-        dim: usize,
-        size: usize,
-        step: usize,
-    ) -> FloatTensor<B>;
+    fn float_unfold(tensor: FloatTensor<B>, dim: usize, size: usize, step: usize)
+    -> FloatTensor<B>;
 }
diff --git a/crates/burn-tensor/src/tests/ops/mod.rs b/crates/burn-tensor/src/tests/ops/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -108,7 +108,12 @@ impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {`
`108`	`108`	`B::bool_repeat_dim(tensor, dim, times)`
`109`	`109`	`}`
`110`	`110`
`111`		`- fn bool_unfold(tensor: BoolTensor<Self>, dim: usize, size: usize, step: usize) -> BoolTensor<Self> {`
	`111`	`+ fn bool_unfold(`
	`112`	`+ tensor: BoolTensor<Self>,`
	`113`	`+ dim: usize,`
	`114`	`+ size: usize,`
	`115`	`+ step: usize,`
	`116`	`+ ) -> BoolTensor<Self> {`
`112`	`117`	`B::bool_unfold(tensor, dim, size, step)`
`113`	`118`	`}`
`114`	`119`	`}`
Original file line number	Diff line number	Diff line change
`@@ -378,7 +378,12 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {`
`378`	`378`	`B::int_cast(tensor, dtype)`
`379`	`379`	`}`
`380`	`380`
`381`		`- fn int_unfold(tensor: IntTensor<Self>, dim: usize, size: usize, step: usize) -> IntTensor<Self> {`
	`381`	`+ fn int_unfold(`
	`382`	`+ tensor: IntTensor<Self>,`
	`383`	`+ dim: usize,`
	`384`	`+ size: usize,`
	`385`	`+ step: usize,`
	`386`	`+ ) -> IntTensor<Self> {`
`382`	`387`	`B::int_unfold(tensor, dim, size, step)`
`383`	`388`	`}`
`384`	`389`	`}`
Original file line number	Diff line number	Diff line change
`@@ -2593,7 +2593,12 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>`
`2593`	`2593`	`// TODO: Implement float_prod and float_sum`
`2594`	`2594`	`// https://github.com/tracel-ai/burn/issues/1458`
`2595`	`2595`
`2596`		`- fn float_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {`
	`2596`	`+ fn float_unfold(`
	`2597`	`+ tensor: FloatTensor<Self>,`
	`2598`	`+ dim: usize,`
	`2599`	`+ size: usize,`
	`2600`	`+ step: usize,`
	`2601`	`+ ) -> FloatTensor<Self> {`
`2597`	`2602`	`AutodiffTensor::new(B::float_unfold(tensor.primitive, dim, size, step))`
`2598`	`2603`	`}`
`2599`	`2604`	`}`
Original file line number	Diff line number	Diff line change
`@@ -142,7 +142,12 @@ where`
`142`	`142`	`kernel::flip::<R, BT, BT>(tensor, axes)`
`143`	`143`	`}`
`144`	`144`
`145`		`- fn bool_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {`
	`145`	`+ fn bool_unfold(`
	`146`	`+ tensor: FloatTensor<Self>,`
	`147`	`+ dim: usize,`
	`148`	`+ size: usize,`
	`149`	`+ step: usize,`
	`150`	`+ ) -> FloatTensor<Self> {`
`146`	`151`	`unfold(tensor, dim, size, step)`
`147`	`152`	`}`
`148`	`153`	`}`
Original file line number	Diff line number	Diff line change
`@@ -703,7 +703,12 @@ where`
`703`	`703`	`}`
`704`	`704`	`}`
`705`	`705`
`706`		`- fn float_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {`
	`706`	`+ fn float_unfold(`
	`707`	`+ tensor: FloatTensor<Self>,`
	`708`	`+ dim: usize,`
	`709`	`+ size: usize,`
	`710`	`+ step: usize,`
	`711`	`+ ) -> FloatTensor<Self> {`
`707`	`712`	`unfold(tensor, dim, size, step)`
`708`	`713`	`}`
`709`	`714`	`}`
Original file line number	Diff line number	Diff line change
`@@ -685,7 +685,12 @@ where`
`685`	`685`	`)`
`686`	`686`	`}`
`687`	`687`
`688`		`- fn int_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {`
	`688`	`+ fn int_unfold(`
	`689`	`+ tensor: FloatTensor<Self>,`
	`690`	`+ dim: usize,`
	`691`	`+ size: usize,`
	`692`	`+ step: usize,`
	`693`	`+ ) -> FloatTensor<Self> {`
`689`	`694`	`unfold(tensor, dim, size, step)`
`690`	`695`	`}`
`691`	`696`	`}`