Clean up

Author: DePasqualeOrg

Libraries/MLXVLM/Models/Gemma3n.swift

@@ -395,12 +395,12 @@ private protocol Gemma3nRMSNormProtocol: UnaryLayer {
 private class Gemma3nRMSNormWithScale: Module, Gemma3nRMSNormProtocol {
     let eps: Float
     let scaleShift: Float
-    @ModuleInfo var weight: MLXArray
+    let weight: MLXArray
 
     init(dim: Int, eps: Float = 1e-6, scaleShift: Float = 0.0) {
         self.eps = eps
         self.scaleShift = scaleShift
-        self._weight.wrappedValue = MLXArray.ones([dim])
+        self.weight = MLXArray.ones([dim])
         super.init()
     }
 
@@ -495,8 +495,8 @@ private class Gemma3nRotaryEmbedding: Module {
     let originalMaxSeqLen: Int
     let config: TextConfig
     let attentionScaling: Float
-    let _invFreq: MLXArray
-    let _originalInvFreq: MLXArray
+    private let _invFreq: MLXArray
+    private let _originalInvFreq: MLXArray
 
     init(config: TextConfig) {
         if let ropeScaling = config.ropeScaling {
@@ -750,7 +750,7 @@ private class Gemma3nAltUp: Module {
     @ModuleInfo(key: "prediction_coefs") var predictionCoefs: Linear
     @ModuleInfo(key: "modality_router") var modalityRouter: Linear
     @ModuleInfo(key: "router_norm") var routerNorm: Gemma3nRMSNormWithScale
-    let _routerInputScale: MLXArray
+    private let _routerInputScale: MLXArray
 
     let config: TextConfig
 
@@ -2340,7 +2340,7 @@ private class Gemma3nAudioAttention: Module {
 
     @ModuleInfo(key: "relative_position_embedding") var relativePositionEmbedding:
         Gemma3nAudioRelativePositionEmbedding
-    @ModuleInfo(key: "per_dim_scale") var perDimScale: MLXArray
+    private let _perDimScale: MLXArray
     @ModuleInfo(key: "q_proj") var qProj: Linear
     @ModuleInfo(key: "k_proj") var kProj: Linear
     @ModuleInfo(key: "v_proj") var vProj: Linear
@@ -2359,7 +2359,7 @@ private class Gemma3nAudioAttention: Module {
 
         self._relativePositionEmbedding.wrappedValue = Gemma3nAudioRelativePositionEmbedding(
             config: config)
-        self._perDimScale.wrappedValue = MLXArray.zeros([headDim])
+        self._perDimScale = MLXArray.zeros([headDim])
 
         self._qProj.wrappedValue = Linear(hiddenSize, numHeads * headDim, bias: false)
         self._kProj.wrappedValue = Linear(hiddenSize, numHeads * headDim, bias: false)
@@ -2460,7 +2460,7 @@ private class Gemma3nAudioAttention: Module {
             Array(x.shape.dropLast()) + [numHeads, headDim]
         )
 
-        let perDimScaleSp = logAddExp(perDimScale, MLXArray(0.0))
+        let perDimScaleSp = logAddExp(_perDimScale, MLXArray(0.0))
         let broadcastShape = [1, 1, 1, headDim]
         let perDimScaleSpBroadcast = perDimScaleSp.reshaped(broadcastShape)
         let scaledQueryStates = queryStates * qScale * perDimScaleSpBroadcast
@@ -2728,19 +2728,19 @@ private class Gemma3nAudioConformerBlock: Module {
 // MARK: - Layer Scale 2D
 private class LayerScale2d: Module, UnaryLayer {
     let inplace: Bool
-    @ModuleInfo var gamma: MLXArray
+    private let _gamma: MLXArray
 
     init(dim: Int, initValues: Float = 1e-5, inplace: Bool = false) {
         self.inplace = inplace
-        self._gamma.wrappedValue = MLXArray(initValues) * MLXArray.ones([dim])
+        self._gamma = MLXArray(initValues) * MLXArray.ones([dim])
         super.init()
     }
 
     func callAsFunction(_ x: MLXArray) -> MLXArray {
         if inplace {
-            return x * gamma
+            return x * _gamma
         } else {
-            return x * gamma
+            return x * _gamma
         }
     }
 }