osobhandClaude Opus 4.6
02d382d5f6
Consistent formatting pass: line wrapping, import sorting, trailing whitespace removal, let-chain indentation, merged derive attributes, and unsafe block reformatting. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
854 lines
27 KiB
Rust
854 lines
27 KiB
Rust
//! Weight loading and management for neural operators.
|
|
//!
|
|
//! This module provides functionality to load pre-trained weights from
|
|
//! SafeTensors format into neural operator models.
|
|
//!
|
|
//! # Supported Formats
|
|
//!
|
|
//! - SafeTensors (`.safetensors`)
|
|
//! - Sharded SafeTensors (for large models)
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_neural_operator::weights::{FNOWeights, load_fno_weights};
|
|
//!
|
|
//! // Load weights from SafeTensors file
|
|
//! let weights = load_fno_weights("./weights/fno2d_darcy/model.safetensors")?;
|
|
//!
|
|
//! // Apply to model
|
|
//! let fno = FNO2d::from_weights(weights, &device)?;
|
|
//! ```
|
|
|
|
use crate::Result;
|
|
use std::collections::HashMap;
|
|
use std::fs::File;
|
|
use std::io::{BufReader, Read, Seek, SeekFrom};
|
|
use std::path::Path;
|
|
|
|
/// Error types for weight loading
|
|
#[derive(Debug, Clone)]
|
|
pub enum WeightError {
|
|
/// File not found
|
|
FileNotFound(String),
|
|
/// Invalid SafeTensors format
|
|
InvalidFormat(String),
|
|
/// Missing required weight
|
|
MissingWeight(String),
|
|
/// Shape mismatch between expected and actual tensor.
|
|
ShapeMismatch {
|
|
/// Name of the tensor with mismatched shape.
|
|
name: String,
|
|
/// Expected shape of the tensor.
|
|
expected: Vec<usize>,
|
|
/// Actual shape found in the file.
|
|
got: Vec<usize>,
|
|
},
|
|
/// I/O error
|
|
IoError(String),
|
|
}
|
|
|
|
impl std::fmt::Display for WeightError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
WeightError::FileNotFound(path) => write!(f, "Weight file not found: {}", path),
|
|
WeightError::InvalidFormat(msg) => write!(f, "Invalid SafeTensors format: {}", msg),
|
|
WeightError::MissingWeight(name) => write!(f, "Missing weight tensor: {}", name),
|
|
WeightError::ShapeMismatch {
|
|
name,
|
|
expected,
|
|
got,
|
|
} => {
|
|
write!(
|
|
f,
|
|
"Shape mismatch for {}: expected {:?}, got {:?}",
|
|
name, expected, got
|
|
)
|
|
}
|
|
WeightError::IoError(msg) => write!(f, "I/O error: {}", msg),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for WeightError {}
|
|
|
|
/// Data type for tensors in SafeTensors format.
|
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
|
pub enum DType {
|
|
/// 16-bit floating point (half precision).
|
|
F16,
|
|
/// 16-bit brain floating point.
|
|
BF16,
|
|
/// 32-bit floating point (single precision).
|
|
F32,
|
|
/// 64-bit floating point (double precision).
|
|
F64,
|
|
/// 8-bit signed integer.
|
|
I8,
|
|
/// 16-bit signed integer.
|
|
I16,
|
|
/// 32-bit signed integer.
|
|
I32,
|
|
/// 64-bit signed integer.
|
|
I64,
|
|
/// 8-bit unsigned integer.
|
|
U8,
|
|
/// Boolean value.
|
|
Bool,
|
|
}
|
|
|
|
impl DType {
|
|
/// Size of each element in bytes
|
|
pub fn size(&self) -> usize {
|
|
match self {
|
|
DType::F16 | DType::BF16 | DType::I16 => 2,
|
|
DType::F32 | DType::I32 => 4,
|
|
DType::F64 | DType::I64 => 8,
|
|
DType::I8 | DType::U8 | DType::Bool => 1,
|
|
}
|
|
}
|
|
|
|
/// Parse from string
|
|
pub fn from_str(s: &str) -> Option<Self> {
|
|
match s.to_uppercase().as_str() {
|
|
"F16" => Some(DType::F16),
|
|
"BF16" => Some(DType::BF16),
|
|
"F32" => Some(DType::F32),
|
|
"F64" => Some(DType::F64),
|
|
"I8" => Some(DType::I8),
|
|
"I16" => Some(DType::I16),
|
|
"I32" => Some(DType::I32),
|
|
"I64" => Some(DType::I64),
|
|
"U8" => Some(DType::U8),
|
|
"BOOL" => Some(DType::Bool),
|
|
_ => None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Metadata for a single tensor in SafeTensors format.
|
|
#[derive(Debug, Clone)]
|
|
pub struct TensorInfo {
|
|
/// Name of the tensor (e.g., "layer.0.weight").
|
|
pub name: String,
|
|
/// Data type of the tensor elements.
|
|
pub dtype: DType,
|
|
/// Shape of the tensor as a list of dimensions.
|
|
pub shape: Vec<usize>,
|
|
/// Byte offsets (start, end) in the data section.
|
|
pub data_offsets: (usize, usize),
|
|
}
|
|
|
|
/// Raw weight data loaded from SafeTensors.
|
|
#[derive(Debug)]
|
|
pub struct RawWeight {
|
|
/// Tensor metadata (name, dtype, shape, offsets).
|
|
pub info: TensorInfo,
|
|
/// Raw byte data of the tensor.
|
|
pub data: Vec<u8>,
|
|
}
|
|
|
|
impl RawWeight {
|
|
/// Convert raw bytes to f32 vector
|
|
pub fn to_f32(&self) -> Result<Vec<f32>> {
|
|
match self.info.dtype {
|
|
DType::F32 => {
|
|
let floats: Vec<f32> = self
|
|
.data
|
|
.chunks_exact(4)
|
|
.map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
|
|
.collect();
|
|
Ok(floats)
|
|
}
|
|
DType::F16 => {
|
|
// Convert f16 to f32
|
|
let floats: Vec<f32> = self
|
|
.data
|
|
.chunks_exact(2)
|
|
.map(|chunk| {
|
|
let bits = u16::from_le_bytes([chunk[0], chunk[1]]);
|
|
half_to_f32(bits)
|
|
})
|
|
.collect();
|
|
Ok(floats)
|
|
}
|
|
DType::BF16 => {
|
|
// Convert bf16 to f32
|
|
let floats: Vec<f32> = self
|
|
.data
|
|
.chunks_exact(2)
|
|
.map(|chunk| {
|
|
let bits = u16::from_le_bytes([chunk[0], chunk[1]]);
|
|
bf16_to_f32(bits)
|
|
})
|
|
.collect();
|
|
Ok(floats)
|
|
}
|
|
_ => Err(crate::NeuralOperatorError::Config(format!(
|
|
"Unsupported dtype for f32 conversion: {:?}",
|
|
self.info.dtype
|
|
))),
|
|
}
|
|
}
|
|
|
|
/// Get the shape of the tensor
|
|
pub fn shape(&self) -> &[usize] {
|
|
&self.info.shape
|
|
}
|
|
|
|
/// Get number of elements
|
|
pub fn numel(&self) -> usize {
|
|
self.info.shape.iter().product()
|
|
}
|
|
}
|
|
|
|
/// Convert f16 bits to f32
|
|
fn half_to_f32(bits: u16) -> f32 {
|
|
let sign = ((bits >> 15) & 1) as u32;
|
|
let exp = ((bits >> 10) & 0x1f) as u32;
|
|
let mantissa = (bits & 0x3ff) as u32;
|
|
|
|
if exp == 0 {
|
|
if mantissa == 0 {
|
|
// Zero
|
|
f32::from_bits(sign << 31)
|
|
} else {
|
|
// Subnormal
|
|
let mut m = mantissa;
|
|
let mut e = 0i32;
|
|
while (m & 0x400) == 0 {
|
|
m <<= 1;
|
|
e -= 1;
|
|
}
|
|
m &= 0x3ff;
|
|
let f32_exp = (127 - 15 + e + 1) as u32;
|
|
let f32_mantissa = m << 13;
|
|
f32::from_bits((sign << 31) | (f32_exp << 23) | f32_mantissa)
|
|
}
|
|
} else if exp == 31 {
|
|
// Inf or NaN
|
|
if mantissa == 0 {
|
|
f32::from_bits((sign << 31) | (0xff << 23))
|
|
} else {
|
|
f32::from_bits((sign << 31) | (0xff << 23) | (mantissa << 13))
|
|
}
|
|
} else {
|
|
// Normal
|
|
let f32_exp = exp + 127 - 15;
|
|
let f32_mantissa = mantissa << 13;
|
|
f32::from_bits((sign << 31) | (f32_exp << 23) | f32_mantissa)
|
|
}
|
|
}
|
|
|
|
/// Convert bf16 bits to f32
|
|
fn bf16_to_f32(bits: u16) -> f32 {
|
|
// BF16 is just the upper 16 bits of f32
|
|
f32::from_bits((bits as u32) << 16)
|
|
}
|
|
|
|
/// Loaded SafeTensors file
|
|
pub struct SafeTensorsFile {
|
|
tensors: HashMap<String, RawWeight>,
|
|
metadata: HashMap<String, String>,
|
|
}
|
|
|
|
impl SafeTensorsFile {
|
|
/// Load SafeTensors file from path
|
|
pub fn load<P: AsRef<Path>>(path: P) -> std::result::Result<Self, WeightError> {
|
|
let path = path.as_ref();
|
|
|
|
let file = File::open(path)
|
|
.map_err(|e| WeightError::FileNotFound(format!("{}: {}", path.display(), e)))?;
|
|
|
|
let mut reader = BufReader::new(file);
|
|
|
|
// Read header size (8 bytes, little-endian u64)
|
|
let mut header_size_bytes = [0u8; 8];
|
|
reader
|
|
.read_exact(&mut header_size_bytes)
|
|
.map_err(|e| WeightError::IoError(e.to_string()))?;
|
|
|
|
let header_size = u64::from_le_bytes(header_size_bytes) as usize;
|
|
|
|
// Read header JSON
|
|
let mut header_bytes = vec![0u8; header_size];
|
|
reader
|
|
.read_exact(&mut header_bytes)
|
|
.map_err(|e| WeightError::IoError(e.to_string()))?;
|
|
|
|
let header_str = String::from_utf8(header_bytes)
|
|
.map_err(|e| WeightError::InvalidFormat(format!("Invalid UTF-8 in header: {}", e)))?;
|
|
|
|
// Parse header JSON
|
|
let header: serde_json::Value = serde_json::from_str(&header_str)
|
|
.map_err(|e| WeightError::InvalidFormat(format!("Invalid JSON in header: {}", e)))?;
|
|
|
|
let header_obj = header
|
|
.as_object()
|
|
.ok_or_else(|| WeightError::InvalidFormat("Header is not a JSON object".to_string()))?;
|
|
|
|
// Parse metadata if present
|
|
let metadata: HashMap<String, String> = if let Some(meta) = header_obj.get("__metadata__") {
|
|
if let Some(obj) = meta.as_object() {
|
|
obj.iter()
|
|
.filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
|
|
.collect()
|
|
} else {
|
|
HashMap::new()
|
|
}
|
|
} else {
|
|
HashMap::new()
|
|
};
|
|
|
|
// Data starts after header
|
|
let data_start = 8 + header_size;
|
|
|
|
// Parse tensor infos
|
|
let mut tensors = HashMap::new();
|
|
|
|
for (name, info) in header_obj {
|
|
if name == "__metadata__" {
|
|
continue;
|
|
}
|
|
|
|
let info_obj = info.as_object().ok_or_else(|| {
|
|
WeightError::InvalidFormat(format!("Tensor {} info is not an object", name))
|
|
})?;
|
|
|
|
let dtype_str = info_obj
|
|
.get("dtype")
|
|
.and_then(|v| v.as_str())
|
|
.ok_or_else(|| {
|
|
WeightError::InvalidFormat(format!("Missing dtype for tensor {}", name))
|
|
})?;
|
|
|
|
let dtype = DType::from_str(dtype_str).ok_or_else(|| {
|
|
WeightError::InvalidFormat(format!(
|
|
"Unknown dtype {} for tensor {}",
|
|
dtype_str, name
|
|
))
|
|
})?;
|
|
|
|
let shape: Vec<usize> = info_obj
|
|
.get("shape")
|
|
.and_then(|v| v.as_array())
|
|
.ok_or_else(|| {
|
|
WeightError::InvalidFormat(format!("Missing shape for tensor {}", name))
|
|
})?
|
|
.iter()
|
|
.filter_map(|v| v.as_u64().map(|n| n as usize))
|
|
.collect();
|
|
|
|
let data_offsets = info_obj
|
|
.get("data_offsets")
|
|
.and_then(|v| v.as_array())
|
|
.ok_or_else(|| {
|
|
WeightError::InvalidFormat(format!("Missing data_offsets for tensor {}", name))
|
|
})?;
|
|
|
|
let start = data_offsets
|
|
.first()
|
|
.and_then(serde_json::Value::as_u64)
|
|
.ok_or_else(|| {
|
|
WeightError::InvalidFormat(format!("Invalid data_offsets for tensor {}", name))
|
|
})? as usize;
|
|
|
|
let end = data_offsets
|
|
.get(1)
|
|
.and_then(serde_json::Value::as_u64)
|
|
.ok_or_else(|| {
|
|
WeightError::InvalidFormat(format!("Invalid data_offsets for tensor {}", name))
|
|
})? as usize;
|
|
|
|
let tensor_info = TensorInfo {
|
|
name: name.clone(),
|
|
dtype,
|
|
shape,
|
|
data_offsets: (start, end),
|
|
};
|
|
|
|
// Read tensor data
|
|
let data_len = end - start;
|
|
let mut data = vec![0u8; data_len];
|
|
|
|
reader
|
|
.seek(SeekFrom::Start((data_start + start) as u64))
|
|
.map_err(|e| WeightError::IoError(e.to_string()))?;
|
|
|
|
reader
|
|
.read_exact(&mut data)
|
|
.map_err(|e| WeightError::IoError(e.to_string()))?;
|
|
|
|
tensors.insert(
|
|
name.clone(),
|
|
RawWeight {
|
|
info: tensor_info,
|
|
data,
|
|
},
|
|
);
|
|
}
|
|
|
|
Ok(Self { tensors, metadata })
|
|
}
|
|
|
|
/// Get list of tensor names
|
|
pub fn tensor_names(&self) -> Vec<&str> {
|
|
self.tensors
|
|
.keys()
|
|
.map(std::string::String::as_str)
|
|
.collect()
|
|
}
|
|
|
|
/// Get tensor by name
|
|
pub fn get(&self, name: &str) -> Option<&RawWeight> {
|
|
self.tensors.get(name)
|
|
}
|
|
|
|
/// Get metadata
|
|
pub fn metadata(&self) -> &HashMap<String, String> {
|
|
&self.metadata
|
|
}
|
|
|
|
/// Check if tensor exists
|
|
pub fn contains(&self, name: &str) -> bool {
|
|
self.tensors.contains_key(name)
|
|
}
|
|
}
|
|
|
|
/// FNO2d pre-trained weights (neuraloperator v2.0 compatible).
|
|
#[derive(Debug)]
|
|
pub struct FNO2dWeights {
|
|
/// Lifting MLP first layer weights: [2*width, in_channels+2].
|
|
pub lifting_fc1_weight: Vec<f32>,
|
|
/// Lifting MLP first layer bias: [2*width].
|
|
pub lifting_fc1_bias: Vec<f32>,
|
|
|
|
/// Lifting MLP second layer weights: [width, 2*width].
|
|
pub lifting_fc2_weight: Vec<f32>,
|
|
/// Lifting MLP second layer bias: [width].
|
|
pub lifting_fc2_bias: Vec<f32>,
|
|
|
|
/// Spectral convolution weights for each layer
|
|
/// Each entry is (weights1_real, weights1_imag, weights2_real, weights2_imag)
|
|
/// Shape: [width, width, modes1, modes2]
|
|
pub spectral_weights: Vec<SpectralConvWeights>,
|
|
|
|
/// 1x1 convolution weights for each layer
|
|
/// Shape: [width, width]
|
|
pub conv_weights: Vec<(Vec<f32>, Vec<f32>)>, // (weight, bias)
|
|
|
|
/// Projection layer weights
|
|
/// projection.0: [128, width], projection.1: [out_channels, 128]
|
|
pub projection_weights: Vec<(Vec<f32>, Vec<f32>)>, // (weight, bias)
|
|
|
|
/// Model configuration
|
|
pub config: FNO2dConfig,
|
|
}
|
|
|
|
/// Spectral convolution weights (complex-valued).
|
|
#[derive(Debug)]
|
|
pub struct SpectralConvWeights {
|
|
/// Real part of first spectral weight matrix.
|
|
pub weights1_real: Vec<f32>,
|
|
/// Imaginary part of first spectral weight matrix.
|
|
pub weights1_imag: Vec<f32>,
|
|
/// Real part of second spectral weight matrix.
|
|
pub weights2_real: Vec<f32>,
|
|
/// Imaginary part of second spectral weight matrix.
|
|
pub weights2_imag: Vec<f32>,
|
|
}
|
|
|
|
/// FNO2d model configuration.
|
|
#[derive(Debug, Clone)]
|
|
pub struct FNO2dConfig {
|
|
/// Number of input channels.
|
|
pub in_channels: usize,
|
|
/// Number of output channels.
|
|
pub out_channels: usize,
|
|
/// Width of hidden channels in Fourier blocks.
|
|
pub width: usize,
|
|
/// Number of Fourier modes to keep (modes1, modes2).
|
|
pub n_modes: (usize, usize),
|
|
/// Number of Fourier blocks.
|
|
pub n_layers: usize,
|
|
/// PDE type this model was trained for (e.g., "darcy", "navier_stokes").
|
|
pub pde_type: String,
|
|
}
|
|
|
|
impl Default for FNO2dConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
in_channels: 3,
|
|
out_channels: 1,
|
|
width: 32,
|
|
n_modes: (12, 12),
|
|
n_layers: 4,
|
|
pde_type: "unknown".to_string(),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Load FNO2d weights from SafeTensors file
|
|
pub fn load_fno2d_weights<P: AsRef<Path>>(
|
|
weights_path: P,
|
|
) -> std::result::Result<FNO2dWeights, WeightError> {
|
|
let safetensors = SafeTensorsFile::load(weights_path)?;
|
|
|
|
// Try to determine config from weights
|
|
let config = infer_fno2d_config(&safetensors)?;
|
|
|
|
// Load lifting MLP layers (neuraloperator v2.0 format)
|
|
// Try new format first: lifting.fcs.0.weight/bias and lifting.fcs.1.weight/bias
|
|
let (lifting_fc1_weight, lifting_fc1_bias, lifting_fc2_weight, lifting_fc2_bias) =
|
|
if safetensors.contains("lifting.fcs.0.weight") {
|
|
// neuraloperator v2.0 format
|
|
let fc1_weight = safetensors
|
|
.get("lifting.fcs.0.weight")
|
|
.ok_or_else(|| WeightError::MissingWeight("lifting.fcs.0.weight".to_string()))?
|
|
.to_f32()
|
|
.map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
let fc1_bias = safetensors.get("lifting.fcs.0.bias").map_or_else(
|
|
|| vec![0.0; 2 * config.width],
|
|
|w| w.to_f32().unwrap_or_default(),
|
|
);
|
|
|
|
let fc2_weight = safetensors
|
|
.get("lifting.fcs.1.weight")
|
|
.ok_or_else(|| WeightError::MissingWeight("lifting.fcs.1.weight".to_string()))?
|
|
.to_f32()
|
|
.map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
let fc2_bias = safetensors.get("lifting.fcs.1.bias").map_or_else(
|
|
|| vec![0.0; config.width],
|
|
|w| w.to_f32().unwrap_or_default(),
|
|
);
|
|
|
|
(fc1_weight, fc1_bias, fc2_weight, fc2_bias)
|
|
} else {
|
|
// Legacy single-layer format - convert to MLP format
|
|
// Create proper 2-layer MLP weights that preserve approximate behavior
|
|
let hidden_dim = 2 * config.width; // 64 for width=32
|
|
|
|
// fc1: [2*width, in_channels+2] - initialize with small random values
|
|
// This will learn to project from in_channels+2 to 2*width
|
|
let fc1_in = config.in_channels + 2; // +2 for positional encoding
|
|
let fc1_weight = vec![0.02; hidden_dim * fc1_in];
|
|
let fc1_bias = vec![0.0; hidden_dim];
|
|
|
|
// fc2: [width, 2*width] - initialize with Xavier-like init
|
|
let fc2_weight = vec![0.02; config.width * hidden_dim];
|
|
let fc2_bias = vec![0.0; config.width];
|
|
|
|
(fc1_weight, fc1_bias, fc2_weight, fc2_bias)
|
|
};
|
|
|
|
// Load spectral convolution weights
|
|
let mut spectral_weights = Vec::with_capacity(config.n_layers);
|
|
for i in 0..config.n_layers {
|
|
let w1r = safetensors
|
|
.get(&format!("spectral_conv.{}.weights1_real", i))
|
|
.ok_or_else(|| {
|
|
WeightError::MissingWeight(format!("spectral_conv.{}.weights1_real", i))
|
|
})?
|
|
.to_f32()
|
|
.map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
let w1i = safetensors
|
|
.get(&format!("spectral_conv.{}.weights1_imag", i))
|
|
.ok_or_else(|| {
|
|
WeightError::MissingWeight(format!("spectral_conv.{}.weights1_imag", i))
|
|
})?
|
|
.to_f32()
|
|
.map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
let w2r = safetensors
|
|
.get(&format!("spectral_conv.{}.weights2_real", i))
|
|
.ok_or_else(|| {
|
|
WeightError::MissingWeight(format!("spectral_conv.{}.weights2_real", i))
|
|
})?
|
|
.to_f32()
|
|
.map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
let w2i = safetensors
|
|
.get(&format!("spectral_conv.{}.weights2_imag", i))
|
|
.ok_or_else(|| {
|
|
WeightError::MissingWeight(format!("spectral_conv.{}.weights2_imag", i))
|
|
})?
|
|
.to_f32()
|
|
.map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
spectral_weights.push(SpectralConvWeights {
|
|
weights1_real: w1r,
|
|
weights1_imag: w1i,
|
|
weights2_real: w2r,
|
|
weights2_imag: w2i,
|
|
});
|
|
}
|
|
|
|
// Load 1x1 conv weights
|
|
let mut conv_weights = Vec::with_capacity(config.n_layers);
|
|
for i in 0..config.n_layers {
|
|
let weight = safetensors
|
|
.get(&format!("conv.{}.weight", i))
|
|
.ok_or_else(|| WeightError::MissingWeight(format!("conv.{}.weight", i)))?
|
|
.to_f32()
|
|
.map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
let bias = safetensors.get(&format!("conv.{}.bias", i)).map_or_else(
|
|
|| vec![0.0; config.width],
|
|
|w| w.to_f32().unwrap_or_default(),
|
|
);
|
|
|
|
conv_weights.push((weight, bias));
|
|
}
|
|
|
|
// Load projection layers
|
|
let mut projection_weights = Vec::new();
|
|
|
|
// First projection layer
|
|
let proj0_weight = safetensors
|
|
.get("projection.0.weight")
|
|
.ok_or_else(|| WeightError::MissingWeight("projection.0.weight".to_string()))?
|
|
.to_f32()
|
|
.map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
let proj0_bias = safetensors
|
|
.get("projection.0.bias")
|
|
.map_or_else(|| vec![0.0; 128], |w| w.to_f32().unwrap_or_default());
|
|
|
|
projection_weights.push((proj0_weight, proj0_bias));
|
|
|
|
// Second projection layer
|
|
let proj1_weight = safetensors
|
|
.get("projection.1.weight")
|
|
.ok_or_else(|| WeightError::MissingWeight("projection.1.weight".to_string()))?
|
|
.to_f32()
|
|
.map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
let proj1_bias = safetensors.get("projection.1.bias").map_or_else(
|
|
|| vec![0.0; config.out_channels],
|
|
|w| w.to_f32().unwrap_or_default(),
|
|
);
|
|
|
|
projection_weights.push((proj1_weight, proj1_bias));
|
|
|
|
Ok(FNO2dWeights {
|
|
lifting_fc1_weight,
|
|
lifting_fc1_bias,
|
|
lifting_fc2_weight,
|
|
lifting_fc2_bias,
|
|
spectral_weights,
|
|
conv_weights,
|
|
projection_weights,
|
|
config,
|
|
})
|
|
}
|
|
|
|
/// Infer FNO2d configuration from weights
|
|
fn infer_fno2d_config(
|
|
safetensors: &SafeTensorsFile,
|
|
) -> std::result::Result<FNO2dConfig, WeightError> {
|
|
// Get width and in_channels from lifting layer
|
|
// Try neuraloperator v2.0 format first (lifting MLP)
|
|
let (width, in_channels) = if safetensors.contains("lifting.fcs.1.weight") {
|
|
// New format: fc2 output is width
|
|
let fc2 = safetensors
|
|
.get("lifting.fcs.1.weight")
|
|
.ok_or_else(|| WeightError::MissingWeight("lifting.fcs.1.weight".to_string()))?;
|
|
let fc2_shape = fc2.shape();
|
|
let width = fc2_shape[0];
|
|
|
|
// fc1 input is in_channels + 2 (with positional encoding)
|
|
let fc1 = safetensors
|
|
.get("lifting.fcs.0.weight")
|
|
.ok_or_else(|| WeightError::MissingWeight("lifting.fcs.0.weight".to_string()))?;
|
|
let fc1_shape = fc1.shape();
|
|
let in_channels = fc1_shape[1].saturating_sub(2); // Remove 2 for positional encoding
|
|
|
|
(width, in_channels)
|
|
} else {
|
|
// Legacy format
|
|
let lifting = safetensors
|
|
.get("lifting.weight")
|
|
.ok_or_else(|| WeightError::MissingWeight("lifting.weight".to_string()))?;
|
|
let lifting_shape = lifting.shape();
|
|
(lifting_shape[0], lifting_shape[1])
|
|
};
|
|
|
|
// Get out_channels from projection layer
|
|
let proj = safetensors
|
|
.get("projection.1.weight")
|
|
.ok_or_else(|| WeightError::MissingWeight("projection.1.weight".to_string()))?;
|
|
|
|
let out_channels = proj.shape()[0];
|
|
|
|
// Count layers by looking for spectral_conv weights
|
|
let mut n_layers = 0;
|
|
while safetensors.contains(&format!("spectral_conv.{}.weights1_real", n_layers)) {
|
|
n_layers += 1;
|
|
}
|
|
|
|
// Get modes from spectral conv shape
|
|
let spectral = safetensors
|
|
.get("spectral_conv.0.weights1_real")
|
|
.ok_or_else(|| WeightError::MissingWeight("spectral_conv.0.weights1_real".to_string()))?;
|
|
|
|
let spectral_shape = spectral.shape();
|
|
let n_modes = if spectral_shape.len() >= 4 {
|
|
(spectral_shape[2], spectral_shape[3])
|
|
} else {
|
|
(12, 12) // Default
|
|
};
|
|
|
|
// Get PDE type from metadata
|
|
let pde_type = safetensors
|
|
.metadata()
|
|
.get("pde_type")
|
|
.cloned()
|
|
.unwrap_or_else(|| "unknown".to_string());
|
|
|
|
Ok(FNO2dConfig {
|
|
in_channels,
|
|
out_channels,
|
|
width,
|
|
n_modes,
|
|
n_layers,
|
|
pde_type,
|
|
})
|
|
}
|
|
|
|
/// Load configuration from JSON file
|
|
pub fn load_config<P: AsRef<Path>>(path: P) -> std::result::Result<FNO2dConfig, WeightError> {
|
|
let file = File::open(path.as_ref()).map_err(|e| WeightError::FileNotFound(e.to_string()))?;
|
|
|
|
let config: serde_json::Value =
|
|
serde_json::from_reader(file).map_err(|e| WeightError::InvalidFormat(e.to_string()))?;
|
|
|
|
Ok(FNO2dConfig {
|
|
in_channels: config["in_channels"].as_u64().unwrap_or(3) as usize,
|
|
out_channels: config["out_channels"].as_u64().unwrap_or(1) as usize,
|
|
width: config["width"].as_u64().unwrap_or(32) as usize,
|
|
n_modes: (
|
|
config["n_modes"]
|
|
.as_array()
|
|
.and_then(|a| a.first())
|
|
.and_then(serde_json::Value::as_u64)
|
|
.unwrap_or(12) as usize,
|
|
config["n_modes"]
|
|
.as_array()
|
|
.and_then(|a| a.get(1))
|
|
.and_then(serde_json::Value::as_u64)
|
|
.unwrap_or(12) as usize,
|
|
),
|
|
n_layers: config["n_layers"].as_u64().unwrap_or(4) as usize,
|
|
pde_type: config["pde_type"].as_str().unwrap_or("unknown").to_string(),
|
|
})
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::io::Write;
|
|
use tempfile::NamedTempFile;
|
|
|
|
fn create_test_safetensors() -> NamedTempFile {
|
|
// Create a minimal SafeTensors file for testing
|
|
let mut file = NamedTempFile::new().unwrap();
|
|
|
|
// Minimal header with one tensor
|
|
let header = r#"{"test":{"dtype":"F32","shape":[2,3],"data_offsets":[0,24]}}"#;
|
|
let header_bytes = header.as_bytes();
|
|
let header_size = header_bytes.len() as u64;
|
|
|
|
// Write header size
|
|
file.write_all(&header_size.to_le_bytes()).unwrap();
|
|
// Write header
|
|
file.write_all(header_bytes).unwrap();
|
|
// Write tensor data (6 f32 values)
|
|
let data: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
|
|
for val in data {
|
|
file.write_all(&val.to_le_bytes()).unwrap();
|
|
}
|
|
|
|
file
|
|
}
|
|
|
|
#[test]
|
|
fn test_safetensors_load() {
|
|
let file = create_test_safetensors();
|
|
let result = SafeTensorsFile::load(file.path());
|
|
assert!(result.is_ok());
|
|
|
|
let st = result.unwrap();
|
|
assert!(st.contains("test"));
|
|
assert_eq!(st.tensor_names().len(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_raw_weight_to_f32() {
|
|
let file = create_test_safetensors();
|
|
let st = SafeTensorsFile::load(file.path()).unwrap();
|
|
|
|
let weight = st.get("test").unwrap();
|
|
let data = weight.to_f32().unwrap();
|
|
|
|
assert_eq!(data.len(), 6);
|
|
assert_eq!(data[0], 1.0);
|
|
assert_eq!(data[5], 6.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dtype_size() {
|
|
assert_eq!(DType::F32.size(), 4);
|
|
assert_eq!(DType::F16.size(), 2);
|
|
assert_eq!(DType::BF16.size(), 2);
|
|
assert_eq!(DType::F64.size(), 8);
|
|
assert_eq!(DType::I8.size(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_dtype_from_str() {
|
|
assert_eq!(DType::from_str("F32"), Some(DType::F32));
|
|
assert_eq!(DType::from_str("f16"), Some(DType::F16));
|
|
assert_eq!(DType::from_str("BF16"), Some(DType::BF16));
|
|
assert_eq!(DType::from_str("invalid"), None);
|
|
}
|
|
|
|
#[test]
|
|
fn test_half_to_f32() {
|
|
// Test zero
|
|
assert_eq!(half_to_f32(0), 0.0);
|
|
|
|
// Test one (0x3C00 in f16)
|
|
let one = half_to_f32(0x3C00);
|
|
assert!((one - 1.0).abs() < 1e-3);
|
|
|
|
// Test negative one (0xBC00 in f16)
|
|
let neg_one = half_to_f32(0xBC00);
|
|
assert!((neg_one + 1.0).abs() < 1e-3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bf16_to_f32() {
|
|
// BF16 is upper 16 bits of f32
|
|
// 1.0 in f32 is 0x3F800000, so in bf16 it's 0x3F80
|
|
let one = bf16_to_f32(0x3F80);
|
|
assert_eq!(one, 1.0);
|
|
|
|
// -1.0 in f32 is 0xBF800000, so in bf16 it's 0xBF80
|
|
let neg_one = bf16_to_f32(0xBF80);
|
|
assert_eq!(neg_one, -1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_fno2d_config_default() {
|
|
let config = FNO2dConfig::default();
|
|
assert_eq!(config.in_channels, 3);
|
|
assert_eq!(config.out_channels, 1);
|
|
assert_eq!(config.width, 32);
|
|
assert_eq!(config.n_modes, (12, 12));
|
|
assert_eq!(config.n_layers, 4);
|
|
}
|
|
}
|