emelex 1.1.1

//! Neural-net building blocks that understand quantized checkpoints.

use std::collections::HashMap;

use crate::engine::{
	array::Array,
	error::{Error, Result},
	ops,
	quant::{QuantParams, Quantization},
};

/// All tensors of a checkpoint plus its quantization config; modules pull
/// their parameters out of this store by canonical path.
pub struct WeightMap {
	tensors: HashMap<String, Array>,
	quant: Quantization,
}

impl WeightMap {
	pub fn new(tensors: HashMap<String, Array>, quant: Quantization) -> Self {
		WeightMap { tensors, quant }
	}

	pub fn quantization(&self) -> &Quantization {
		&self.quant
	}

	pub fn contains(&self, key: &str) -> bool {
		self.tensors.contains_key(key)
	}

	pub fn keys(&self) -> impl Iterator<Item = &String> {
		self.tensors.keys()
	}

	/// Borrow a tensor without consuming it (read-only shape/dtype access).
	///
	/// emelex patch (not upstream): non-consuming inspection API for MTP
	/// detection. All pre-sanitize detection (raw-HF predicate, prefix probe,
	/// sentinel/gate-shape validation) uses `peek`
	/// exclusively — detection is contractually non-mutating (no
	/// take/reinsert while probing).
	pub fn peek(&self, key: &str) -> Option<&Array> {
		self.tensors.get(key)
	}

	/// Remove and return a raw tensor, erroring with the path when missing.
	pub fn take(&mut self, key: &str) -> Result<Array> {
		self.tensors
			.remove(key)
			.ok_or_else(|| Error::Model(format!("missing weight: {key}")))
	}

	pub fn take_optional(&mut self, key: &str) -> Option<Array> {
		self.tensors.remove(key)
	}

	/// Insert or overwrite a tensor (used by `sanitize` steps that
	/// synthesize new keys, e.g. splitting/stacking fused MoE weights).
	pub fn insert(&mut self, key: String, value: Array) {
		self.tensors.insert(key, value);
	}

	/// Rename keys with a mapping function (used by model `sanitize` steps).
	pub fn rename_keys(&mut self, f: impl Fn(&str) -> Option<String>) {
		let keys: Vec<String> = self.tensors.keys().cloned().collect();
		for key in keys {
			if let Some(new_key) = f(&key) {
				if new_key != key {
					if let Some(value) = self.tensors.remove(&key) {
						self.tensors.insert(new_key, value);
					}
				}
			} else {
				self.tensors.remove(&key);
			}
		}
	}

	/// Apply the same normalization to per-layer quantization override keys.
	pub fn normalize_quant_keys(&mut self, f: impl Fn(&str) -> Option<String>) {
		let entries: Vec<(String, crate::engine::quant::LayerOverride)> =
			self.quant.per_layer.drain().collect();
		for (key, value) in entries {
			let new_key = f(&key).unwrap_or(key);
			self.quant.per_layer.insert(new_key, value);
		}
	}

	/// Load a linear layer at `path`, choosing dense vs quantized from the
	/// checkpoint contents and config (mlx-lm `class_predicate` semantics).
	pub fn linear(&mut self, path: &str) -> Result<Linear> {
		let has_scales = self.contains(&format!("{path}.scales"));
		let params = self.quant.resolve(path, has_scales);
		let weight = self.take(&format!("{path}.weight"))?;
		let bias = self.take_optional(&format!("{path}.bias"));

		match params {
			Some(q) if has_scales => {
				let scales = self.take(&format!("{path}.scales"))?;
				let biases = self.take_optional(&format!("{path}.biases"));
				Ok(Linear::Quantized(QuantizedLinear {
					weight,
					scales,
					biases,
					params: q,
					bias,
				}))
			}
			_ => Ok(Linear::Dense(DenseLinear { weight, bias })),
		}
	}

	/// Load an embedding table (dense or quantized).
	pub fn embedding(&mut self, path: &str) -> Result<Embedding> {
		let has_scales = self.contains(&format!("{path}.scales"));
		let params = self.quant.resolve(path, has_scales);
		let weight = self.take(&format!("{path}.weight"))?;

		match params {
			Some(q) if has_scales => {
				let scales = self.take(&format!("{path}.scales"))?;
				let biases = self.take_optional(&format!("{path}.biases"));
				Ok(Embedding::Quantized {
					weight,
					scales,
					biases,
					params: q,
				})
			}
			_ => Ok(Embedding::Dense { weight }),
		}
	}

	/// Load an RMSNorm weight vector.
	pub fn rms_norm(&mut self, path: &str, eps: f32) -> Result<RmsNorm> {
		Ok(RmsNorm {
			weight: self.take(&format!("{path}.weight"))?,
			eps,
		})
	}

	/// Load a (weight, bias) LayerNorm.
	pub fn layer_norm(&mut self, path: &str, eps: f32) -> Result<LayerNorm> {
		Ok(LayerNorm {
			weight: self.take(&format!("{path}.weight"))?,
			bias: self.take_optional(&format!("{path}.bias")),
			eps,
		})
	}
}

/// Dense (float) linear layer.
pub struct DenseLinear {
	pub weight: Array,
	pub bias: Option<Array>,
}

/// Quantized linear layer backed by fused `quantized_matmul`.
pub struct QuantizedLinear {
	pub weight: Array,
	pub scales: Array,
	pub biases: Option<Array>,
	pub params: QuantParams,
	pub bias: Option<Array>,
}

/// A linear layer as stored in the checkpoint: dense or quantized.
pub enum Linear {
	Dense(DenseLinear),
	Quantized(QuantizedLinear),
}

impl Linear {
	pub fn forward(&self, x: &Array) -> Result<Array> {
		let out = match self {
			Linear::Dense(l) => {
				let wt = ops::swapaxes(&l.weight, -1, -2)?;
				ops::matmul(x, &wt)?
			}
			Linear::Quantized(l) => ops::quantized_matmul(
				x,
				&l.weight,
				&l.scales,
				l.biases.as_ref(),
				true,
				l.params.group_size,
				l.params.bits,
				l.params.mode,
			)?,
		};
		match self.bias() {
			Some(b) => ops::add(&out, b),
			None => Ok(out),
		}
	}

	fn bias(&self) -> Option<&Array> {
		match self {
			Linear::Dense(l) => l.bias.as_ref(),
			Linear::Quantized(l) => l.bias.as_ref(),
		}
	}

	/// Output feature count (rows of the stored weight matrix).
	pub fn output_dims(&self) -> i32 {
		match self {
			Linear::Dense(l) => l.weight.dim(0),
			Linear::Quantized(l) => l.weight.dim(0),
		}
	}

	/// The dtype an input must be cast to before [`Linear::forward`] (only
	/// meaningful for [`Linear::Dense`]; quantized weights accept any
	/// floating input dtype via `quantized_matmul`, so this returns the
	/// input dtype unchanged in that case).
	pub fn weight_dtype(
		&self,
		input_dtype: crate::engine::array::Dtype,
	) -> crate::engine::array::Dtype {
		match self {
			Linear::Dense(l) => l.weight.dtype(),
			Linear::Quantized(_) => input_dtype,
		}
	}
}

/// Token embedding table, optionally quantized, usable as a tied LM head.
pub enum Embedding {
	Dense {
		weight: Array,
	},
	Quantized {
		weight: Array,
		scales: Array,
		biases: Option<Array>,
		params: QuantParams,
	},
}

impl Embedding {
	/// Look up embeddings for integer token ids.
	pub fn forward(&self, ids: &Array) -> Result<Array> {
		match self {
			Embedding::Dense { weight } => ops::take_axis(weight, ids, 0),
			Embedding::Quantized {
				weight,
				scales,
				biases,
				params,
			} => {
				// Gather packed rows then dequantize just those rows.
				let w = ops::take_axis(weight, ids, 0)?;
				let s = ops::take_axis(scales, ids, 0)?;
				let b = match biases {
					Some(b) => Some(ops::take_axis(b, ids, 0)?),
					None => None,
				};
				ops::dequantize(
					&w,
					&s,
					b.as_ref(),
					params.group_size,
					params.bits,
					params.mode,
				)
			}
		}
	}

	/// Use the (transposed) embedding table as the LM head projection.
	pub fn as_linear(&self, x: &Array) -> Result<Array> {
		match self {
			Embedding::Dense { weight } => {
				let wt = ops::swapaxes(weight, -1, -2)?;
				ops::matmul(x, &wt)
			}
			Embedding::Quantized {
				weight,
				scales,
				biases,
				params,
			} => ops::quantized_matmul(
				x,
				weight,
				scales,
				biases.as_ref(),
				true,
				params.group_size,
				params.bits,
				params.mode,
			),
		}
	}
}

/// RMS normalization with a learned scale.
pub struct RmsNorm {
	pub weight: Array,
	pub eps: f32,
}

impl RmsNorm {
	pub fn forward(&self, x: &Array) -> Result<Array> {
		ops::rms_norm(x, Some(&self.weight), self.eps)
	}
}

/// Standard (mean/variance) layer normalization with an affine weight and
/// optional bias.
pub struct LayerNorm {
	pub weight: Array,
	pub bias: Option<Array>,
	pub eps: f32,
}

impl LayerNorm {
	pub fn forward(&self, x: &Array) -> Result<Array> {
		ops::layer_norm(x, Some(&self.weight), self.bias.as_ref(), self.eps)
	}
}