Skip to content

Make ~ in the mask DSL actually negate - #93

Open
cemde wants to merge 1 commit into
liukidar:mainfrom
cemde:fix/68-mask-negation
Open

Make ~ in the mask DSL actually negate#93
cemde wants to merge 1 commit into
liukidar:mainfrom
cemde:fix/68-mask-negation

Conversation

@cemde

@cemde cemde commented Aug 9, 2026

Copy link
Copy Markdown
Collaborator

Bug

M has two distinct entry points, and they do different jobs:

  • apply(leaf) is the per-leaf predicate. Given one Param, return True or False.
  • __call__(pydag) is the tree walker. It reffs the tree, evaluates apply at every Param, and returns a copy with rejected leaves replaced by None.

~M(A) constructs a _M_not, which needs to invert the predicate. It overrode __call__ instead of apply, so the negation was installed on the walker and never reached the predicate.

Impact

Two silent failures.

(~M(A))(model) stopped walking the tree. It evaluated the predicate once against the whole model and returned a bare True rather than a masked pytree.

Worse, _M_and.apply and _M_or.apply resolve their operands by calling .apply on each one. _M_not never overrode apply, so it inherited M's positive version. M(A) & ~M(B) therefore computed M(A) & M(B), selecting exactly the parameters the user asked to exclude, which trains the wrong parameters and freezes the ones meant to learn. That expression appears in M's own class docstring. De Morgan's laws did not hold either.

Fix

_M_not now defines apply instead of __call__:

class _M_not(M):
    def apply(self, leaf: Any):
        return not M._resolve(self.mask, leaf)

That matches _M_or, _M_and and _M_hasattr, which all override apply and none of which touch __call__. The negation now sits where M._resolve looks for it, so the combinators see it, and _M_not inherits the working M.__call__ for tree application.

Also in this PR

Reviewing this branch surfaced a separate defect. M.__call__ calls tree_ref first, which replaces duplicate Param references with _BaseParamRef placeholders holding integer indices. A positive mask rejects those because they are not the requested type; a negated mask selects them for the same reason, and Optim.apply_updates then adds the raw index to the weight. On a shared model, one step gives 2.8 where -0.2 is correct.

Filed as #88 and guarded by a bug-marked test added here. Not fixed: it predates this change (M(None) selects them on main today), but ~ working is what makes the path reachable in practice.

Closes #68

_M_not defined __call__, the tree-application entry point, instead of
apply, the per-leaf predicate. Negation was therefore never applied by
the &/| combinators and (~M(A))(model) returned a bare bool.
@cemde
cemde force-pushed the fix/68-mask-negation branch from d25dde5 to 9e80b21 Compare August 9, 2026 18:23
@cemde
cemde requested a review from liukidar August 9, 2026 18:34
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.

~ in the mask DSL never negates

1 participant