torch

crutcher · crutcher · commit 91fdc4df8b76 · 2025-09-22T18:29:31.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
@@ -2563,7 +2563,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
@@ -127,7 +127,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
@@ -684,7 +684,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
@@ -662,7 +662,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, TensorData, TensorMetadata,
     ops::{BoolTensor, BoolTensorOps, FloatTensor, IntTensor, binary_ops_shape},
 };
+use std::cmp::max;
 use std::marker::PhantomData;
 
 use crate::{
@@ -747,7 +751,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,
@@ -757,11 +766,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, 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> {
@@ -2265,7 +2265,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,
@@ -2275,11 +2280,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
@@ -2178,7 +2178,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,
@@ -2188,11 +2193,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
@@ -179,10 +179,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
@@ -419,6 +419,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
@@ -2138,7 +2138,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
`@@ -2563,7 +2563,12 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>`
`2563`	`2563`	`// TODO: Implement float_prod and float_sum`
`2564`	`2564`	`// https://github.com/tracel-ai/burn/issues/1458`
`2565`	`2565`
`2566`		`- fn float_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {`
	`2566`	`+ fn float_unfold(`
	`2567`	`+ tensor: FloatTensor<Self>,`
	`2568`	`+ dim: usize,`
	`2569`	`+ size: usize,`
	`2570`	`+ step: usize,`
	`2571`	`+ ) -> FloatTensor<Self> {`
`2567`	`2572`	`AutodiffTensor::new(B::float_unfold(tensor.primitive, dim, size, step))`
`2568`	`2573`	`}`
`2569`	`2574`	`}`
Original file line number	Diff line number	Diff line change
`@@ -127,7 +127,12 @@ where`
`127`	`127`	`kernel::flip::<R, BT, BT>(tensor, axes)`
`128`	`128`	`}`
`129`	`129`
`130`		`- fn bool_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {`
	`130`	`+ fn bool_unfold(`
	`131`	`+ tensor: FloatTensor<Self>,`
	`132`	`+ dim: usize,`
	`133`	`+ size: usize,`
	`134`	`+ step: usize,`
	`135`	`+ ) -> FloatTensor<Self> {`
`131`	`136`	`unfold(tensor, dim, size, step)`
`132`	`137`	`}`
`133`	`138`	`}`
Original file line number	Diff line number	Diff line change
`@@ -684,7 +684,12 @@ where`
`684`	`684`	`}`
`685`	`685`	`}`
`686`	`686`
`687`		`- fn float_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {`
	`687`	`+ fn float_unfold(`
	`688`	`+ tensor: FloatTensor<Self>,`
	`689`	`+ dim: usize,`
	`690`	`+ size: usize,`
	`691`	`+ step: usize,`
	`692`	`+ ) -> FloatTensor<Self> {`
`688`	`693`	`unfold(tensor, dim, size, step)`
`689`	`694`	`}`
`690`	`695`	`}`
Original file line number	Diff line number	Diff line change
`@@ -662,7 +662,12 @@ where`
`662`	`662`	`)`
`663`	`663`	`}`
`664`	`664`
`665`		`- fn int_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {`
	`665`	`+ fn int_unfold(`
	`666`	`+ tensor: FloatTensor<Self>,`
	`667`	`+ dim: usize,`
	`668`	`+ size: usize,`
	`669`	`+ step: usize,`
	`670`	`+ ) -> FloatTensor<Self> {`
`666`	`671`	`unfold(tensor, dim, size, step)`
`667`	`672`	`}`
`668`	`673`	`}`