burn_tensor/tensor/api/
int.rs

1use crate::{
2    Float, Int, Shape, Tensor, TensorData, TensorPrimitive, backend::Backend, cartesian_grid,
3};
4
5use core::ops::Range;
6
7impl<B> Tensor<B, 1, Int>
8where
9    B: Backend,
10{
11    /// Returns a new integer tensor on the specified device.
12    ///
13    /// # Arguments
14    ///
15    /// * `range` - The range of values to generate.
16    /// * `device` - The device to create the tensor on.
17    pub fn arange(range: Range<i64>, device: &B::Device) -> Self {
18        Tensor::new(B::int_arange(range, device))
19    }
20
21    /// Returns a new integer tensor on the specified device.
22    ///
23    /// # Arguments
24    ///
25    /// * `range` - The range of values to generate.
26    /// * `step` - The step between each value.
27    pub fn arange_step(range: Range<i64>, step: usize, device: &B::Device) -> Self {
28        Tensor::new(B::int_arange_step(range, step, device))
29    }
30}
31
32impl<const D: usize, B> Tensor<B, D, Int>
33where
34    B: Backend,
35{
36    /// Create a tensor from integers (i32), placing it on a given device.
37    ///
38    /// # Example
39    ///
40    /// ```rust
41    /// use burn_tensor::backend::Backend;
42    /// use burn_tensor::{Tensor, Int};
43    ///
44    /// fn example<B: Backend>() {
45    ///     let device = B::Device::default();
46    ///     let _x: Tensor<B, 1, Int> = Tensor::from_ints([1, 2], &device);
47    ///     let _y: Tensor<B, 2, Int> = Tensor::from_ints([[1, 2], [3, 4]], &device);
48    /// }
49    /// ```
50    pub fn from_ints<A: Into<TensorData>>(ints: A, device: &B::Device) -> Self {
51        Self::from_data(ints.into().convert::<i32>(), device)
52    }
53
54    /// Returns a new tensor with the same shape and device as the current tensor and the data
55    /// cast to Float.
56    ///
57    /// # Example
58    ///
59    /// ```rust
60    /// use burn_tensor::backend::Backend;
61    /// use burn_tensor::{Int, Tensor};
62    ///
63    /// fn example<B: Backend>() {
64    ///     let device = Default::default();
65    ///     let int_tensor = Tensor::<B, 1, Int>::arange(0..5, &device);
66    ///     let float_tensor = int_tensor.float();
67    /// }
68    /// ```
69    pub fn float(self) -> Tensor<B, D, Float> {
70        Tensor::new(TensorPrimitive::Float(B::int_into_float(self.primitive)))
71    }
72
73    /// Generates a cartesian grid for the given tensor shape on the specified device.
74    /// The generated tensor is of dimension `D2 = D + 1`, where each element at dimension D contains the cartesian grid coordinates for that element.
75    ///
76    /// # Arguments
77    ///
78    /// * `shape` - The shape specifying the dimensions of the tensor.
79    /// * `device` - The device to create the tensor on.
80    ///
81    /// # Panics
82    ///
83    /// Panics if `D2` is not equal to `D+1`.
84    ///
85    /// # Examples
86    ///
87    /// ```rust
88    ///    use burn_tensor::Int;
89    ///    use burn_tensor::{backend::Backend, Shape, Tensor};
90    ///    fn example<B: Backend>() {
91    ///        let device = Default::default();
92    ///        let result: Tensor<B, 3, _> = Tensor::<B, 2, Int>::cartesian_grid([2, 3], &device);
93    ///        println!("{}", result);
94    ///    }
95    /// ```
96    pub fn cartesian_grid<S: Into<Shape>, const D2: usize>(
97        shape: S,
98        device: &B::Device,
99    ) -> Tensor<B, D2, Int> {
100        cartesian_grid::<B, S, D, D2>(shape, device)
101    }
102
103    /// Applies the bitwise logical and operation with each bit representing the integer.
104    pub fn bitwise_and(self, other: Self) -> Self {
105        Self::new(B::bitwise_and(self.primitive, other.primitive))
106    }
107
108    /// Applies the bitwise logical or operation with another tensor.
109    pub fn bitwise_or(self, other: Self) -> Self {
110        Self::new(B::bitwise_or(self.primitive, other.primitive))
111    }
112
113    /// Applies the bitwise logical xor operation with another tensor.
114    pub fn bitwise_xor(self, other: Self) -> Self {
115        Self::new(B::bitwise_xor(self.primitive, other.primitive))
116    }
117
118    /// Applies the bitwise logical not operation.
119    pub fn bitwise_not(self) -> Self {
120        Self::new(B::bitwise_not(self.primitive))
121    }
122
123    /// Applies the bitwise logical and operation with each bit in the scalar and the integers in the tensor.
124    pub fn bitwise_and_scalar(self, other: B::IntElem) -> Self {
125        Self::new(B::bitwise_and_scalar(self.primitive, other))
126    }
127
128    /// Applies the bitwise logical or operation with each bit in the scalar and the integers in the tensor.
129    pub fn bitwise_or_scalar(self, other: B::IntElem) -> Self {
130        Self::new(B::bitwise_or_scalar(self.primitive, other))
131    }
132
133    /// Applies bitwise logical xor operation with each bit in the scalar and the integers in the tensor.
134    pub fn bitwise_xor_scalar(self, other: B::IntElem) -> Self {
135        Self::new(B::bitwise_xor_scalar(self.primitive, other))
136    }
137
138    /// Applies the bitwise left shift operation with the integers in the tensor.
139    pub fn bitwise_left_shift(self, other: Self) -> Self {
140        Self::new(B::bitwise_left_shift(self.primitive, other.primitive))
141    }
142
143    /// Applies the bitwise right shift operation with the integers in the tensor.
144    pub fn bitwise_right_shift(self, other: Self) -> Self {
145        Self::new(B::bitwise_right_shift(self.primitive, other.primitive))
146    }
147
148    /// Applies the bitwise left shift operation with the scalar.
149    pub fn bitwise_left_shift_scalar(self, other: B::IntElem) -> Self {
150        Self::new(B::bitwise_left_shift_scalar(self.primitive, other))
151    }
152
153    /// Applies the bitwise right shift operation with the scalar.
154    pub fn bitwise_right_shift_scalar(self, other: B::IntElem) -> Self {
155        Self::new(B::bitwise_right_shift_scalar(self.primitive, other))
156    }
157}
burn_tensor/tensor/api/int.rs

burn_tensor/tensor/api/
int.rs
