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

burn_tensor/tensor/api/
int.rs
