25 lines
723 B
Swift
25 lines
723 B
Swift
// Copyright © 2024 Apple Inc.
|
|
|
|
import MLXNN
|
|
|
|
extension Module {
|
|
|
|
/// Compute the number of parameters in a possibly quantized model
|
|
public func numParameters() -> Int {
|
|
return leafModules().flattenedValues().map {
|
|
mod -> Int in
|
|
if let qlin = mod as? QuantizedLinear {
|
|
return qlin.scales.size * qlin.groupSize
|
|
} else if let qemb = mod as? QuantizedEmbedding {
|
|
return qemb.scales.size * qemb.groupSize
|
|
} else {
|
|
return mod.parameters().flattenedValues().reduce(
|
|
0,
|
|
{
|
|
$0 + $1.size
|
|
})
|
|
}
|
|
}.reduce(0, +)
|
|
}
|
|
}
|