Validate freeze and unfreeze keys against the whole model when recursing

[ayaangazali]

Proposed changes

freeze and unfreeze document keys="bias" as a whole model operation ("freeze all biases by calling module.freeze(keys="bias")", and unfreeze's example is nn.Transformer() then model.unfreeze(keys="bias")). But with strict=True they validate the keys separately inside every visited submodule, so the first submodule without a bias aborts the whole call. That makes the two documented features unusable together on any model with mixed layers, including the bundled Transformer:

import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_flatten

m = nn.Transformer(dims=8, num_heads=2, num_encoder_layers=1, num_decoder_layers=0)
len([k for k, _ in tree_flatten(m.parameters()) if k.endswith("bias")])   # 6

m.freeze(keys="bias", strict=True)
# KeyError: "Module doesn't contain member bias."

It happens because MultiHeadAttention's projections default to bias=False while the MLP Linear layers default to bias=True, so the traversal hits a module with no bias and raises even though six of them exist.

When recursing, this now checks the keys once against the whole model and only raises if a key is present nowhere. Behavior that stays the same:

• a key that really is absent everywhere still raises KeyError, including when it is one entry in a list
• recurse=False still validates against just that module
• strict=False and the no-keys case are untouched

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

(the new test fails on main with the KeyError and passes here; test_nn.py is green)

If you would rather read strict as per module and keep the current behavior, then the docstrings are what need changing instead, and I am glad to flip this into a docs only PR.

---

about me: freshman contributor, and i do use Claude Code while hunting for these, but i traced this one to the bias=False projections in MultiHeadAttention myself and checked the recurse and strict combinations by hand before writing it up. tell me if the whole model reading is not the one you intended.

[zcbenz]

Strictly speaking this is a breaking change, but I agree that the feature is not useful at all with current state.