markbaseengine/Sources/MarkBase/Weights/TensorDescriptor.swift

/// Metadata for a single tensor stored in a SafeTensors file.
public struct TensorDescriptor: Sendable, Codable {
    public let name: String
    public let dtype: TensorDType
    public let shape: [Int]
    /// Byte offset from the start of the safetensors data section.
    public let dataOffset: Int
    /// Byte size of the tensor data.
    public let dataSize: Int

    /// Total number of elements.
    public var elementCount: Int { shape.reduce(1, *) }

    /// Check if shape is compatible with a given dim count.
    public func hasRank(_ rank: Int) -> Bool { shape.count == rank }

    /// For quantized tensors: returns the grouping factor (elements per group).
    /// MLX default: 64 elements per quantization group (for Gemma 4 E4B 4-bit).
    public var quantizationGroupSize: Int { 64 }
}

/// Group of tensors that together represent a quantized linear layer.
///   weight: U32 packed (shape: [outDim, inDim / 32 * 4])
///   scales: BF16    (shape: [outDim, inDim / 32])
///   biases: BF16    (shape: [outDim, inDim / 32])
public struct QuantizedTensorGroup: Sendable {
    public let name: String
    public let weight: TensorDescriptor
    public let scales: TensorDescriptor
    public let biases: TensorDescriptor

    /// Output dimension.
    public var outDim: Int { weight.shape[0] }
    /// Input dimension (pre-quantization).
    public var inDim: Int { scales.shape[1] * 32 }
    /// Block size (elements per group).
    public let groupSize: Int = 64
}