entrenar/efficiency/device/tpu.rs
1//! TPU information and detection.
2
3use serde::{Deserialize, Serialize};
4
5/// TPU information
6#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
7pub struct TpuInfo {
8 /// TPU version (e.g., "v4", "v5e")
9 pub version: String,
10 /// Number of TPU cores
11 pub cores: u32,
12 /// High bandwidth memory in bytes
13 pub hbm_bytes: u64,
14}
15
16impl TpuInfo {
17 /// Create new TPU info
18 pub fn new(version: impl Into<String>, cores: u32, hbm_bytes: u64) -> Self {
19 Self { version: version.into(), cores, hbm_bytes }
20 }
21
22 /// Get HBM in GB
23 pub fn hbm_gb(&self) -> f64 {
24 self.hbm_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
25 }
26}