Skip to content

Fix Log and Equal is_equivalent ignoring primitive state - #4266

Draft
kapellirohith wants to merge 1 commit into
ml-explore:mainfrom
kapellirohith:hunt-aug15-probe
Draft

Fix Log and Equal is_equivalent ignoring primitive state#4266
kapellirohith wants to merge 1 commit into
ml-explore:mainfrom
kapellirohith:hunt-aug15-probe

Conversation

@kapellirohith

Copy link
Copy Markdown
Contributor

Problem

Log carries base_ for log / log2 / log10, and Equal carries equal_nan_ for array_equal, but both use DEFINE_DEFAULT_IS_EQUIVALENT(), which returns true unconditionally (mlx/primitives.h:35). The compile simplify pass merges two nodes when array_equivalent accepts them (mlx/compile.cpp:606, applied at :698), and that check delegates the final decision to is_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:

import mlx.core as mx

def entropies(p):
    nats = -mx.sum(p * mx.log(p))
    bits = -mx.sum(p * mx.log2(p))
    return mx.stack([nats, bits])

p = mx.array([0.1, 0.2, 0.3, 0.4])
print(entropies(p))
print(mx.compile(entropies)(p))

Before, the entropy in bits silently returns the value in nats:

array([1.27985, 1.84644], dtype=float32)
array([1.27985, 1.27985], dtype=float32)

After, both lines agree:

array([1.27985, 1.84644], dtype=float32)
array([1.27985, 1.84644], dtype=float32)

Same on cpu and gpu, since the merge happens on the graph. mx.array_equal(x, x) and mx.array_equal(x, x, equal_nan=True) collapse the same way, as do isclose and allclose with and without equal_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_equivalent ignoring width_. These two are the only remaining primitives that declare state() 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 two Log nodes 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 on DEVICE=cpu the only failure is the pre-existing test_fft_too_large, which fails identically on main. A MLX_METAL_JIT=ON build matches, with the pre-existing 291 test_quantized errors unchanged. The new tests ran 200 times per device without a flake, and pre-commit run --all-files is clean.

Simplification is not weakened: same-base logs and same-equal_nan comparisons 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 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)

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.
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.

1 participant