Skip to content

Conversation

petrukha-ivan
Copy link
Contributor

What

  • Fixed setting a value for an optional child module
  • Fixed resetting an optional child module back to nil

Motivation

Consider the following example with an optional child module:

class Example: Module {
    @ModuleInfo(key: "a") var a: Linear = Linear(1, 1)
    @ModuleInfo(key: "b") var b: Linear?
}

Currently, attempting to update b via update(modules:verify:) fails:

let linear = Linear(1, 1)
let update = ModuleChildren(values: ["b": .value(linear)])
try example.update(modules: update, verify: .noUnusedKeys) // Fatal error: Unable to set b on Example(b=nil)

The reason is it has double optional value inside:

(lldb) po (value, type(of: value))
▿ 2 elements
  - .0 : nil
  - .1 : Swift.Optional<Swift.Optional<MLXNN.Linear>>

Comparing to a non-optional child module:

(lldb) po (value, type(of: value))
▿ 2 elements
  ▿ .0 : Optional<Linear>
    ▿ some : Linear(inputDimensions=1, outputDimensions=1, bias=true)
  - .1 : Swift.Optional<MLXNN.Linear>

Due to the double optional, ModuleValue.build(value:) maps the value to .other(nil) instead of .none and prevents updates from being applied correctly. This fix ensures that setting a value when the current value is nil works as expected and resetting an optional module back to nil is supported as well.

Comment on lines +1272 to +1273
case let v as Module? where v == nil:
return .none
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This part handles the double-optional nil case. When the value is already set and not nil, it's handled by the previous case case let v as Module:

Comment on lines +690 to +694
case (.none, .value(let newModule)):
try self.updateModule(key: key, newModule)

case (.value(.module), .none):
try self.updateModule(key: key, Optional<Module>.none as Any)
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This allows setting new module when current value is nil and also resetting it back to nil if needed

Comment on lines -597 to -599
if case .none = value {
return
}
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm a bit concerned about this deletion. While all unit tests pass, there may be edge cases I'm not aware of. But this deletion is needed to fall into switch statement.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this looks ok since there is a .none added to the switch below

@@ -906,4 +906,43 @@ class ModuleTests: XCTestCase {
// this should not throw because _freqs is not considered
try rope.update(parameters: .init(), verify: .all)
}

func testOptionalModuleUpdates() throws {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

great!

Copy link
Collaborator

@davidkoski davidkoski left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice fix and thanks for the tests!

@davidkoski davidkoski merged commit b74974b into ml-explore:main Jun 15, 2025
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants