Skip to content

Check pytree structure in compiled tree assertions - #481

Open
AryanRG wants to merge 1 commit into
google-deepmind:mainfrom
AryanRG:fix/w160-chex
Open

AryanRG wants to merge 1 commit into
google-deepmind:mainfrom
AryanRG:fix/w160-chex

Conversation

@AryanRG

@AryanRG AryanRG commented Oct 27, 2026

Copy link
Copy Markdown

Compiled assert_trees_all_equal and assert_trees_all_close can accept different pytree structures when their flattened leaves match, or raise IndexError when leaves are missing. Check tree definitions before comparing leaves and report a normal assertion failure on a mismatch.

Regressions cover eager and compiled comparisons, structural mismatches, matching trees, custom messages, and disabled assertions. Affected assertion suites pass on CPU.

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