out-of-band PR feedback

- turn off training mode
- clarify documentation
- handle MaterializedArray passed as inputs (state) to a compiled function
	- contents are constant by definition so can be discarded and captured normally

Author: davidkoski

Source/MLX/MaterializedArray.swift

@@ -19,7 +19,7 @@ import Numerics
 /// - You can only create an instance via ``MLXArray/materialized()``
 ///   or ``materialize(_:)->MaterializedArray``
 /// - Mutation methods are marked as unavailable and will `fatalError`
-///   if you somehow manage to call them.
+///   if you somehow manage to call them (e.g. through dynamic types)
 /// - It is declared `@unchecked Sendable` and may be passed freely between
 ///   tasks, actors, and other concurrency boundaries.
 ///

Source/MLX/Transforms+Compile.swift

@@ -43,7 +43,11 @@ final class CompiledFunction: @unchecked (Sendable) {
     }
 
     func innerCall(_ arguments: [MLXArray]) -> [MLXArray] {
-        let stateInputs = inputs.flatMap { $0.innerState() }
+        // MaterializedArray can never change, so it doesn't need to be
+        // threaded through the tracer swap below -- it is simply captured
+        // as a constant, the same as any other array reachable from the
+        // closure that isn't part of `inputs`/`outputs`.
+        let stateInputs = inputs.flatMap { $0.innerState() }.filter { !($0 is MaterializedArray) }
         let argumentsCount = arguments.count
 
         // inner function to hande the compilation.  this is called
@@ -70,7 +74,9 @@ final class CompiledFunction: @unchecked (Sendable) {
             let result = f(tracerArguments)
 
             // recapture the state as it may have changed
-            let stateOutputTracers = outputs.flatMap { $0.innerState() }.map { $0.copyContext() }
+            let stateOutputTracers = outputs.flatMap { $0.innerState() }
+                .filter { !($0 is MaterializedArray) }
+                .map { $0.copyContext() }
 
             // put the original values back in the state
             for (s, saved) in zip(stateInputs, savedStateInputs) {
@@ -119,7 +125,7 @@ final class CompiledFunction: @unchecked (Sendable) {
         let resultsPlusStateOutput = mlx_vector_array_values(resultVector)
 
         // push the stateOutput into the state
-        let stateOutput = outputs.flatMap { $0.innerState() }
+        let stateOutput = outputs.flatMap { $0.innerState() }.filter { !($0 is MaterializedArray) }
 
         for (s, newValues) in zip(stateOutput, resultsPlusStateOutput.suffix(stateOutput.count)) {
             s._updateInternal(newValues)

Source/MLXNN/MaterializedModule.swift

@@ -12,6 +12,8 @@ import MLX
 /// mutation surface so the wrapped module cannot be modified through this
 /// reference.
 ///
+/// This also sets ``Module/training`` to `false`.
+///
 /// Note: only parameters that can be be mutated, e.g. are wrapped with `@ParameterInfo`,
 /// will actually be updated to the `MaterializedArray` type.  All others will
 /// be evaluated and are materialized, even if they do not have the type
@@ -42,17 +44,20 @@ import MLX
 ///
 /// ## What is sealed
 ///
-/// The following `Module` operations are marked `@available(*, unavailable)`
-/// on `MaterializedModule` and will trap if called:
-///
-/// - `update(parameters:...)` and `update(modules:...)`
-/// - `updateModule(key:_:)`
-/// - `apply(filter:map:)`
-/// - `freeze(...)` / `unfreeze(...)`
-/// - `train(_:)`
+/// `MaterializedModule` is not a `Module` and does not provide access to the
+/// wrapped `Module`.  No access to bare `MLXArray` is provided, though
+/// for purposes of introspection the ``parameters()`` can give
+/// `MaterializedArray`.  There are also properties like ``parameterNBytes``
+/// and ``parameterCount`` for callers that just need size information.
 ///
 /// ## Calling the wrapped module
 ///
+/// `MaterializedModule`is not a `Module` subclass and cannot conform
+/// to protocols that have `Module` requirements.  Users may want to split their
+/// protocols into requirements for performing inference (which
+/// `MaterializedModule` is meant for) and for doing training where `Module`
+/// is required and this type is not appropriate (immutable).
+///
 /// `MaterializedModule` does not itself know how to invoke `base`; that is
 /// added per-layer-shape via an extension that constrains `LayerType`.
 /// For example, every ``UnaryLayer`` already supports being called with a
@@ -137,6 +142,8 @@ open class MaterializedModule<LayerType: Module>: IndentedDescription, @unchecke
 
     public init(_ base: consuming LayerType) {
         self._base = base
+
+        self._base.train(false)
         self._base.materialize()
 
         // seal the consumed base so that any retained reference held by a

Source/MLXNN/Module.swift

@@ -1532,6 +1532,7 @@ public enum ModuleValue {
         // cannot check via unwapProperty -- see wrappedValue.set
     }
 
+    // See also ModuleInfo.Setter
     struct Setter: TypeErasedSetter {
         unowned var info: ParameterInfo<T>
 
@@ -1679,6 +1680,7 @@ private protocol TypeErasedSetterProvider {
         }
     }
 
+    // See also ParameterInfo.Setter
     struct Setter: TypeErasedSetter {
         unowned var info: ModuleInfo<T>
 

Source/MLXNN/Quantized.swift

@@ -362,7 +362,11 @@ open class QuantizedLinear: Linear, Quantized {
     }
 
     open override var parameterCount: Int {
-        scales.size * groupSize
+        if biases != nil {
+            scales.size * groupSize * 2
+        } else {
+            scales.size * groupSize
+        }
     }
 
     public override func unfreeze(

Tests/MLXTests/MaterializedTests.swift

@@ -40,6 +40,44 @@ struct MaterializedTests {
         _ = await t.result
     }
 
+    @Test
+    func testCompileInputsOutputsMaterialized() {
+        // A MaterializedArray can never change, so it has nothing for
+        // compile(inputs:outputs:) to observe or update -- it is simply
+        // captured as a constant when the function is traced, the same as
+        // any other value reachable from the closure that isn't part of
+        // `inputs`/`outputs`.  This is the documented
+        // `compile(inputs: [state], outputs: [state], f)` pattern (see
+        // compilation.md).
+        let state = MLXRandom.normal([4, 4]).materialized()
+        let compiled = compile(inputs: [state], outputs: [state]) { (x: MLXArray) in
+            x + state
+        }
+
+        let x = MLXRandom.normal([4, 4])
+        let r1 = compiled(x)
+        let r2 = x + state
+        #expect(r1.allClose(r2).item(Bool.self))
+
+        // same for a `Module` whose parameters were materialized (e.g. via
+        // `MaterializedModule`) -- this is the documented
+        // `compile(inputs: [model], outputs: [model])` pattern.
+        //
+        // Note: capturing the interior module is documented as forbidden but
+        // used here to test.  Do not copy this pattern!
+        let l = Linear(4, 4)
+        _ = MaterializedModule(l)
+
+        let compiledModel = compile(inputs: [l], outputs: [l]) { (x: MLXArray) in
+            l(x)
+        }
+
+        let input = MLXRandom.normal([4, 4])
+        let modelResult = compiledModel(input)
+        let expected = l(input)
+        #expect(modelResult.allClose(expected).item(Bool.self))
+    }
+
     @Test
     func testMaterializedLinear() async {
         let l = Linear(10, 10)