Fix Log and Equal is_equivalent ignoring primitive state - #4266
Draft
kapellirohith wants to merge 1 commit into
Draft
Fix Log and Equal is_equivalent ignoring primitive state#4266kapellirohith wants to merge 1 commit into
kapellirohith wants to merge 1 commit into
Conversation
Log carries base_ and Equal carries equal_nan_, but both used DEFINE_DEFAULT_IS_EQUIVALENT(), so the compile simplify pass merged log/log2/log10 of the same input, and array_equal with and without equal_nan, into a single node.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Logcarriesbase_forlog/log2/log10, andEqualcarriesequal_nan_forarray_equal, but both useDEFINE_DEFAULT_IS_EQUIVALENT(), which returnstrueunconditionally (mlx/primitives.h:35). The compile simplify pass merges two nodes whenarray_equivalentaccepts them (mlx/compile.cpp:606, applied at :698), and that check delegates the final decision tois_equivalent(mlx/compile.cpp:629). Two logs of different bases over the same input therefore look identical and one of them is dropped.Apple M3 Pro, macOS 26.6.1 (25G76). Before is main at 140faa8, after is this branch at 9bfa102:
Before, the entropy in bits silently returns the value in nats:
After, both lines agree:
Same on cpu and gpu, since the merge happens on the graph.
mx.array_equal(x, x)andmx.array_equal(x, x, equal_nan=True)collapse the same way, as doiscloseandallclosewith and withoutequal_nan.Present since 8ca7f9e (2023-11-29). Existing tests miss it because simplify exempts function outputs from merging (mlx/compile.cpp:699), so returning the two logs directly is correct and only intermediates are affected.
Same class as #2978, which fixed
RandomBits::is_equivalentignoringwidth_. These two are the only remaining primitives that declarestate()while using the default equivalence.Testing
Both new tests fail on main and pass with the change. On main the C++ case reports
CHECK( 4335042088 != 4335042088 ), the twoLognodes having been merged into one id.With the change: C++ suite 277/277 on gpu and on
DEVICE=cpu; python suite 841 passing on gpu, and onDEVICE=cputhe only failure is the pre-existingtest_fft_too_large, which fails identically on main. AMLX_METAL_JIT=ONbuild matches, with the pre-existing 291test_quantizederrors unchanged. The new tests ran 200 times per device without a flake, andpre-commit run --all-filesis clean.Simplification is not weakened: same-base logs and same-
equal_nancomparisons still merge, asserted in the C++ test, and compiled graph node counts are unchanged on five functions covering log-softmax, cross entropy, layer norm, RMS norm and attention.Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes