burn_tensor/tensor/api/
int.rs

1use burn_backend::{ElementConversion, Scalar, ops::IntTensorOps};
2use burn_dispatch::Dispatch;
3
4use crate::{
5    Cast, Device, Float, Int, Shape, Tensor, TensorCreationOptions, TensorData, cartesian_grid,
6    ops::BridgeTensor,
7};
8
9use core::ops::Range;
10
11impl Tensor<1, Int> {
12    /// Returns a new integer tensor on the specified device.
13    ///
14    /// # Arguments
15    ///
16    /// * `range` - The range of values to generate.
17    /// * `device` - The device to create the tensor on.
18    pub fn arange(range: Range<i64>, options: impl Into<TensorCreationOptions>) -> Self {
19        let opt = options.into();
20        let dtype = opt.resolve_dtype::<Int>();
21        Tensor::new(arange_impl(range, opt.device, dtype))
22    }
23
24    /// Returns a new integer tensor on the specified device.
25    ///
26    /// # Arguments
27    ///
28    /// * `range` - The range of values to generate.
29    /// * `step` - The step between each value.
30    pub fn arange_step(
31        range: Range<i64>,
32        step: usize,
33        options: impl Into<TensorCreationOptions>,
34    ) -> Self {
35        let opt = options.into();
36        let dtype = opt.resolve_dtype::<Int>();
37        Tensor::new(arange_step_impl(range, step, opt.device, dtype))
38    }
39}
40
41impl<const D: usize> Tensor<D, Int> {
42    /// Create a tensor from integers (i32), placing it on a given device.
43    ///
44    /// # Example
45    ///
46    /// ```rust
47    /// use burn_tensor::{Tensor, Int};
48    ///
49    /// fn example() {
50    ///     let device = Default::default();
51    ///     let _x: Tensor<1, Int> = Tensor::from_ints([1, 2], &device);
52    ///     let _y: Tensor<2, Int> = Tensor::from_ints([[1, 2], [3, 4]], &device);
53    /// }
54    /// ```
55    pub fn from_ints<A: Into<TensorData>>(ints: A, device: &Device) -> Self {
56        Self::from_data(ints.into().convert::<i32>(), device)
57    }
58
59    /// Returns a new tensor with the same shape and device as the current tensor and the data
60    /// cast to Float.
61    ///
62    /// # Example
63    ///
64    /// ```rust
65    /// use burn_tensor::{Int, Tensor};
66    ///
67    /// fn example() {
68    ///     let device = Default::default();
69    ///     let int_tensor = Tensor::<1, Int>::arange(0..5, &device);
70    ///     let float_tensor = int_tensor.float();
71    /// }
72    /// ```
73    pub fn float(self) -> Tensor<D, Float> {
74        let device = self.device();
75        Tensor::new(int_to_float_impl(self.primitive, device))
76    }
77
78    /// Generates a cartesian grid for the given tensor shape on the specified device.
79    /// The generated tensor is of dimension `D2 = D + 1`, where each element at dimension D contains the cartesian grid coordinates for that element.
80    ///
81    /// # Arguments
82    ///
83    /// * `shape` - The shape specifying the dimensions of the tensor.
84    /// * `device` - The device to create the tensor on.
85    ///
86    /// # Panics
87    ///
88    /// Panics if `D2` is not equal to `D+1`.
89    ///
90    /// # Examples
91    ///
92    /// ```rust
93    ///    use burn_tensor::Int;
94    ///    use burn_tensor::{Shape, Tensor};
95    ///    fn example() {
96    ///        let device = Default::default();
97    ///        let result: Tensor<3, _> = Tensor::<2, Int>::cartesian_grid([2, 3], &device);
98    ///        println!("{}", result);
99    ///    }
100    /// ```
101    pub fn cartesian_grid<S: Into<Shape>, const D2: usize>(
102        shape: S,
103        device: &Device,
104    ) -> Tensor<D2, Int> {
105        cartesian_grid::<S, D, D2>(shape, device)
106    }
107
108    /// Applies element wise square operation.
109    ///
110    #[cfg_attr(doc, doc = r#"$y_i = x_i * x_i$"#)]
111    #[cfg_attr(not(doc), doc = "`y_i = x_i * x_i`")]
112    pub fn square(self) -> Self {
113        Self::new(square_impl(self.primitive))
114    }
115
116    /// Applies the bitwise logical and operation with each bit representing the integer.
117    pub fn bitwise_and(self, other: Self) -> Self {
118        Self::new(bitwise_and_impl(self.primitive, other.primitive))
119    }
120
121    /// Applies the bitwise logical or operation with another tensor.
122    pub fn bitwise_or(self, other: Self) -> Self {
123        Self::new(bitwise_or_impl(self.primitive, other.primitive))
124    }
125
126    /// Applies the bitwise logical xor operation with another tensor.
127    pub fn bitwise_xor(self, other: Self) -> Self {
128        Self::new(bitwise_xor_impl(self.primitive, other.primitive))
129    }
130
131    /// Applies the bitwise logical not operation.
132    pub fn bitwise_not(self) -> Self {
133        Self::new(bitwise_not_impl(self.primitive))
134    }
135
136    /// Applies the bitwise logical and operation with each bit in the scalar and the integers in the tensor.
137    pub fn bitwise_and_scalar(self, other: impl ElementConversion) -> Self {
138        let other = Scalar::new(other, &self.dtype());
139        Self::new(bitwise_and_scalar_impl(self.primitive, other))
140    }
141
142    /// Applies the bitwise logical or operation with each bit in the scalar and the integers in the tensor.
143    pub fn bitwise_or_scalar(self, other: impl ElementConversion) -> Self {
144        let other = Scalar::new(other, &self.dtype());
145        Self::new(bitwise_or_scalar_impl(self.primitive, other))
146    }
147
148    /// Applies bitwise logical xor operation with each bit in the scalar and the integers in the tensor.
149    pub fn bitwise_xor_scalar(self, other: impl ElementConversion) -> Self {
150        let other = Scalar::new(other, &self.dtype());
151        Self::new(bitwise_xor_scalar_impl(self.primitive, other))
152    }
153
154    /// Applies the bitwise left shift operation with the integers in the tensor.
155    pub fn bitwise_left_shift(self, other: Self) -> Self {
156        Self::new(bitwise_left_shift_impl(self.primitive, other.primitive))
157    }
158
159    /// Applies the bitwise right shift operation with the integers in the tensor.
160    pub fn bitwise_right_shift(self, other: Self) -> Self {
161        Self::new(bitwise_right_shift_impl(self.primitive, other.primitive))
162    }
163
164    /// Applies the bitwise left shift operation with the scalar.
165    pub fn bitwise_left_shift_scalar(self, other: impl ElementConversion) -> Self {
166        let other = Scalar::new(other, &self.dtype());
167        Self::new(bitwise_left_shift_scalar_impl(self.primitive, other))
168    }
169
170    /// Applies the bitwise right shift operation with the scalar.
171    pub fn bitwise_right_shift_scalar(self, other: impl ElementConversion) -> Self {
172        let other = Scalar::new(other, &self.dtype());
173        Self::new(bitwise_right_shift_scalar_impl(self.primitive, other))
174    }
175
176    /// Converts a tensor to the specified data type.
177    ///
178    /// Supports both within-kind casting (e.g., `IntDType::I64`) and cross-kind casting
179    /// (e.g., `FloatDType::F32` to produce a float tensor).
180    ///
181    /// This is a no-op when casting to the current dtype within the same kind.
182    ///
183    /// # Example
184    ///
185    /// ```rust
186    /// use burn_tensor::{Tensor, Int, IntDType, FloatDType};
187    ///
188    /// fn example() {
189    ///     let device = Default::default();
190    ///     let int_tensor = Tensor::<1, Int>::arange(0..5, &device);
191    ///
192    ///     // Within-kind cast (int to int)
193    ///     let i64_tensor = int_tensor.clone().cast(IntDType::I64);
194    ///
195    ///     // Cross-kind cast (int to float)
196    ///     let float_tensor = int_tensor.cast(FloatDType::F32);
197    /// }
198    /// ```
199    #[must_use]
200    pub fn cast<T: Cast<D, Int>>(self, dtype: T) -> Tensor<D, T::OutputKind> {
201        T::cast(self, dtype)
202    }
203}
204
205// =========================================================================
206// Non-generic implementation helpers (outlined from the generic API).
207// See the crate-level docs for the rationale behind this pattern.
208// =========================================================================
209
210fn arange_impl(range: Range<i64>, device: Device, dtype: burn_std::DType) -> BridgeTensor {
211    BridgeTensor::int(Dispatch::int_arange(
212        range,
213        device.as_dispatch(),
214        dtype.into(),
215    ))
216}
217
218fn arange_step_impl(
219    range: Range<i64>,
220    step: usize,
221    device: Device,
222    dtype: burn_std::DType,
223) -> BridgeTensor {
224    BridgeTensor::int(Dispatch::int_arange_step(
225        range,
226        step,
227        device.as_dispatch(),
228        dtype.into(),
229    ))
230}
231
232fn int_to_float_impl(p: BridgeTensor, device: Device) -> BridgeTensor {
233    let out_dtype = device.settings().float_dtype;
234    BridgeTensor::float(Dispatch::int_into_float(p.into(), out_dtype))
235}
236fn square_impl(tensor: BridgeTensor) -> BridgeTensor {
237    BridgeTensor::int(Dispatch::int_square(tensor.into()))
238}
239fn bitwise_and_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor {
240    BridgeTensor::int(Dispatch::bitwise_and(lhs.into(), rhs.into()))
241}
242fn bitwise_or_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor {
243    BridgeTensor::int(Dispatch::bitwise_or(lhs.into(), rhs.into()))
244}
245fn bitwise_xor_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor {
246    BridgeTensor::int(Dispatch::bitwise_xor(lhs.into(), rhs.into()))
247}
248fn bitwise_not_impl(p: BridgeTensor) -> BridgeTensor {
249    BridgeTensor::int(Dispatch::bitwise_not(p.into()))
250}
251fn bitwise_and_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor {
252    BridgeTensor::int(Dispatch::bitwise_and_scalar(p.into(), other))
253}
254fn bitwise_or_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor {
255    BridgeTensor::int(Dispatch::bitwise_or_scalar(p.into(), other))
256}
257fn bitwise_xor_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor {
258    BridgeTensor::int(Dispatch::bitwise_xor_scalar(p.into(), other))
259}
260fn bitwise_left_shift_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor {
261    BridgeTensor::int(Dispatch::bitwise_left_shift(lhs.into(), rhs.into()))
262}
263fn bitwise_right_shift_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor {
264    BridgeTensor::int(Dispatch::bitwise_right_shift(lhs.into(), rhs.into()))
265}
266fn bitwise_left_shift_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor {
267    BridgeTensor::int(Dispatch::bitwise_left_shift_scalar(p.into(), other))
268}
269fn bitwise_right_shift_scalar_impl(p: BridgeTensor, other: Scalar) -> BridgeTensor {
270    BridgeTensor::int(Dispatch::bitwise_right_shift_scalar(p.into(), other))
271}
burn_tensor/tensor/api/int.rs

burn_tensor/tensor/api/
int.rs
