burn_tensor/tensor/api/
options.rs

1use burn_std::DType;
2
3use crate::{Device, bridge::BasicOps, ops::Kind};
4
5/// Options for tensor creation.
6///
7/// This struct allows specifying the `device` and overriding the data type when creating a tensor.
8/// When the `dtype` is not specified, the [device's default settings](crate::DeviceSettings) is used.
9#[derive(Debug, Clone)]
10pub struct TensorCreationOptions {
11    /// Device where the tensor will be created.
12    pub device: Device,
13    /// Optional data type.
14    /// If `None`, the dtype will be inferred on creation from the [device settings](crate::DeviceSettings).
15    pub dtype: Option<DType>,
16}
17
18impl Default for TensorCreationOptions {
19    /// Returns new options with the backend's default device.
20    fn default() -> Self {
21        Self::new(Default::default())
22    }
23}
24
25impl TensorCreationOptions {
26    /// Create new options with a specific device.
27    ///
28    /// Data type will follow the [device settings](crate::DeviceSettings) on tensor creation.
29    pub fn new(device: Device) -> Self {
30        Self {
31            device,
32            dtype: None,
33        }
34    }
35
36    /// Set the tensor creation data type.
37    pub fn with_dtype(mut self, dtype: DType) -> Self {
38        self.dtype = Some(dtype);
39
40        self
41    }
42
43    /// Set the tensor creation device.
44    pub fn with_device(mut self, device: Device) -> Self {
45        self.device = device;
46
47        self
48    }
49
50    /// Returns the tensor data type, or a provided default if not set.
51    ///
52    /// This is useful for cases where [`TensorCreationOptions`] may not have an explicit `dtype`.
53    pub fn dtype_or(&self, dtype: DType) -> DType {
54        self.dtype.unwrap_or(dtype)
55    }
56
57    /// Returns the tensor data type, or the default from the [device settings](crate::DeviceSettings).
58    pub(crate) fn resolve_dtype<K: BasicOps>(&self) -> DType {
59        let settings = self.device.settings();
60        let default = match K::KIND {
61            Kind::Float => settings.float_dtype.into(),
62            Kind::Int => settings.int_dtype.into(),
63            Kind::Bool => settings.bool_dtype.into(),
64        };
65        match K::KIND {
66            // `BoolStore` variants are backend-specific storage layouts rather than
67            // semantic types, and a requested variant (e.g. forced from a serialized
68            // snapshot by `burn-store`) may be unsupported on the target device, so
69            // always resolve bool tensors to the device's default bool storage.
70            Kind::Bool => default,
71            _ => self.dtype.unwrap_or(default),
72        }
73    }
74}
75
76impl From<&Device> for TensorCreationOptions {
77    /// Convenience conversion from a reference to a device.
78    ///
79    /// Example:
80    /// ```rust
81    /// use burn_tensor::TensorCreationOptions;
82    /// use burn_tensor::Device;
83    ///
84    /// fn example(device: Device) {
85    ///     let options: TensorCreationOptions = (&device).into();
86    /// }
87    /// ```
88    fn from(device: &Device) -> Self {
89        TensorCreationOptions::new(device.clone())
90    }
91}
92
93impl From<(&crate::Device, DType)> for TensorCreationOptions {
94    /// Convenience conversion for a specified `(&device, dtype)` tuple.
95    fn from(args: (&crate::Device, DType)) -> Self {
96        TensorCreationOptions::new(args.0.clone()).with_dtype(args.1)
97    }
98}
burn_tensor/tensor/api/options.rs

burn_tensor/tensor/api/
options.rs
