pub struct Device { /* private fields */ }Expand description
Device representation for tensor operations
A device specifies where tensors are located and where operations should be performed. Each device has a type (CPU or CUDA) and an index (0 for CPU, GPU ID for CUDA). The device system provides thread-safe context management and automatic resource cleanup.
§Fields
device_type- The type of device (CPU or CUDA)index- Device index (0 for CPU, GPU ID for CUDA)
§Examples
use train_station::Device;
// Create CPU device
let cpu = Device::cpu();
assert!(cpu.is_cpu());
assert_eq!(cpu.index(), 0);
assert_eq!(cpu.to_string(), "cpu");
// Create CUDA device (when feature enabled)
#[cfg(feature = "cuda")]
{
if train_station::cuda_is_available() {
let cuda = Device::cuda(0);
assert!(cuda.is_cuda());
assert_eq!(cuda.index(), 0);
assert_eq!(cuda.to_string(), "cuda:0");
}
}§Thread Safety
This type is thread-safe and can be shared between threads. Device contexts are managed per-thread using thread-local storage.
§Memory Layout
The device struct is small and efficient:
- Size: 16 bytes (8 bytes for enum + 8 bytes for index)
- Alignment: 8 bytes
- Copy semantics: Implements Copy for efficient passing
Implementations§
Source§impl Device
impl Device
Sourcepub fn cpu() -> Self
pub fn cpu() -> Self
Create a CPU device
CPU devices always have index 0 and are always available regardless of feature flags or system configuration.
§Returns
A Device representing the CPU (always index 0)
§Examples
use train_station::Device;
let device = Device::cpu();
assert!(device.is_cpu());
assert_eq!(device.index(), 0);
assert_eq!(device.device_type(), train_station::DeviceType::Cpu);Examples found in repository?
179fn demonstrate_utility_functions() {
180 println!("\n--- Utility Functions ---");
181
182 let tensor = Tensor::from_slice(&[1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
183
184 // Basic properties
185 println!("Shape: {:?}", tensor.shape().dims);
186 println!("Size: {}", tensor.size());
187 println!("Is contiguous: {}", tensor.is_contiguous());
188 println!("Device: {:?}", tensor.device());
189
190 // Mathematical operations
191 let sum = tensor.sum();
192 println!("Sum: {}", sum.value());
193
194 let mean = tensor.mean();
195 println!("Mean: {}", mean.value());
196
197 let norm = tensor.norm();
198 println!("Norm: {}", norm.value());
199
200 // Device placement
201 let cpu_tensor = Tensor::zeros_on_device(vec![3, 3], train_station::Device::cpu());
202 println!(
203 "CPU tensor: shape {:?}, device: {:?}",
204 cpu_tensor.shape().dims,
205 cpu_tensor.device()
206 );
207}Sourcepub fn cuda(index: usize) -> Self
pub fn cuda(index: usize) -> Self
Create a CUDA device
Creates a device representing a specific CUDA GPU. The device index must be valid for the current system configuration.
§Arguments
index- CUDA device index (0-based)
§Returns
A Device representing the specified CUDA device
§Panics
Panics in the following cases:
- CUDA feature is not enabled (
--features cudanot specified) - CUDA is not available on the system
- Device index is out of range (>= number of available devices)
§Examples
use train_station::Device;
// CPU device is always available
let cpu = Device::cpu();
// CUDA device (when feature enabled and available)
#[cfg(feature = "cuda")]
{
if train_station::cuda_is_available() {
let device_count = train_station::cuda_device_count();
if device_count > 0 {
let cuda = Device::cuda(0);
assert!(cuda.is_cuda());
assert_eq!(cuda.index(), 0);
}
}
}Sourcepub fn device_type(&self) -> DeviceType
pub fn device_type(&self) -> DeviceType
Trait Implementations§
Source§impl Display for Device
impl Display for Device
Source§fn fmt(&self, f: &mut Formatter<'_>) -> Result
fn fmt(&self, f: &mut Formatter<'_>) -> Result
Format the device as a string
§Returns
String representation of the device:
"cpu"for CPU devices"cuda:{index}"for CUDA devices
§Examples
use train_station::Device;
let cpu = Device::cpu();
assert_eq!(cpu.to_string(), "cpu");
#[cfg(feature = "cuda")]
{
if train_station::cuda_is_available() {
let cuda = Device::cuda(0);
assert_eq!(cuda.to_string(), "cuda:0");
}
}Source§impl From<DeviceType> for Device
impl From<DeviceType> for Device
Source§fn from(device_type: DeviceType) -> Self
fn from(device_type: DeviceType) -> Self
Convert DeviceType to Device with index 0
§Arguments
device_type- The device type to convert
§Returns
A Device with the specified type and index 0
§Panics
Panics if device_type is DeviceType::Cuda and CUDA is not available
or the feature is not enabled.
§Examples
use train_station::{Device, DeviceType};
let cpu_type = DeviceType::Cpu;
let cpu_device = Device::from(cpu_type);
assert!(cpu_device.is_cpu());
assert_eq!(cpu_device.index(), 0);