Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@
never a real invariant on Metal (GPT Q4 confirms); the W1-A parity test
now asserts winners identical + log_scores within tests/conftest.py's
PARITY_ATOL = 1e-2 (measured Metal batch-shape drift ~0.004 nats, winners
stable; coder1's PR #21 imports the same constant post-merge; exact
stable; PR #21 imports the same constant post-merge; exact
equality still holds on the deterministic FakeModel path). trie.py:
score_trie returns (log_probs, legal_mass_logs) — one function, the
legal_mass_at_node callback optional (None = mass 1.0, for the MLX-free
Expand Down
24 changes: 10 additions & 14 deletions jevmlx/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -2915,15 +2915,15 @@ def reconcile_case_constraints(
state: AssembledState,
constraints,
schema: StructuredSchema,
compiled_constraints: "CompiledConstraints | None" = None,
compiled_constraints: CompiledConstraints,
) -> AssembledState:
"""Stage 4 (W5b-10 C1): constrained MAP over the first-pass decisions.

Consumes CompiledConstraints (W5b-11, coder4 — compiled by
validate_constraints_for_schema/compile_constraints) and re-picks the
joint assignment maximizing summed per-field log scores subject to the
case-level constraints. Telemetry is updated in place for changed
fields (value + probability from the winning score key).
Consumes the request's CompiledConstraints (W5b-11 — compiled once by
compile_constraints) and re-picks the joint assignment maximizing
summed per-field log scores subject to the case-level constraints.
Telemetry is updated in place for changed fields (value + probability
from the winning score key).

Returns the updated AssembledState carrying the SAME rescored_fields
plus the reconciled names (parsed_json / field_telemetry are the SAME
Expand All @@ -2934,10 +2934,6 @@ def reconcile_case_constraints(
return AssembledState(
state.parsed_json, state.field_telemetry, state.rescored_fields, tuple()
)
if compiled_constraints is None:
from jevmlx.constraints import compile_constraints

compiled_constraints = compile_constraints(constraints, schema)
parsed_json = state.parsed_json
field_telemetry = state.field_telemetry
field_log_scores = {
Expand Down Expand Up @@ -2987,8 +2983,8 @@ def run_dependency_waves(
A NAMED boundary over _selective_second_pass (the wave loop stays there):
takes the post-MAP AssembledState, returns the updated state (parsed /
telemetry mutated in place by the waves) plus the second-pass telemetry
dict. coder4 wires jevmlx.timing.Ledger's 'dependency' span around this
call — it is the only dependency-stage boundary.
dict. jevmlx.timing.Ledger's 'dependency' span wraps this call — it is
the only dependency-stage boundary.
"""
if not any(f.depends_on is not None for f in schema.fields.values()):
return state, {"rerun_fields": [], "rerun_rows": 0, "second_pass_ms": 0.0}
Expand Down Expand Up @@ -3228,8 +3224,8 @@ def _assemble(
dispatch_rows (row-kind dispatch) -> per-field score_scalar_field /
score_multi_field (each ending in the shared scalar finalizer) ->
reconcile_case_constraints (W3-D MAP over CompiledConstraints) ->
run_dependency_waves (W3-D part 2, the named boundary coder4 wraps in
timing.Ledger's 'dependency' span) -> finalize_public_result (the
run_dependency_waves (W3-D part 2, the named boundary that
timing.Ledger's 'dependency' span wraps) -> finalize_public_result (the
result dict). Everything AFTER the forward passes lives in the stages;
the batched path reuses this unchanged.
"""
Expand Down
4 changes: 2 additions & 2 deletions jevmlx/timing.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,15 @@

Standalone module — pure Python, NO mlx import, NO engine wiring yet.
Engine adoption (deleting the t_gather_ms/t_broadcast_ms/t_scored_ms
accumulators) lands after W5-B with coder2 informed.
accumulators) lands after W5-B.

Why: the current ``*_ms`` keys overlap (``suffix_eval_ms`` includes
broadcast + gather, which are ALSO reported separately), ``elapsed_ms``
excludes the prior pass, and decide_many divides group timers per context.
One ledger measures each interval ONCE, non-overlapping; every reported
key is a derivation of the same interval set.

Phases (per coder6's note): ``prior`` (the neutral pass) and ``main``
Phases: ``prior`` (the neutral pass) and ``main``
(everything else). Names: plan, prompt_render, prefill, cache_merge,
transformer, lm_head, gather, rescore, dependency, reconciliation,
assembly — plus the batched wrappers group_wall / per-context assembly.
Expand Down
2 changes: 1 addition & 1 deletion tests/test_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -471,7 +471,7 @@ def test_w1a_scoring_parity_batch_vs_chunked_real_model(engine):
W3-C's measurement on this machine: since W2-B shortened the slot rows
to 4 tokens the observed worst drift is ~0.029 nats on the fintech_fraud
preset (winner stable) — hence the shared PARITY_ATOL constant in
conftest.py, coordinated with coder3's W2-D tolerance change. The
conftest.py (the shared W2-D tolerance). The
invariant that MATTERS is the decision: the same winner per field, and
log_scores that agree to within FP tolerance. Exact equality is still
asserted on the FakeModel path (test_engine_fake.py) where the model is
Expand Down
2 changes: 1 addition & 1 deletion tests/test_timing.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ def test_close_ordering_nested():

def test_derived_flat_keys_match_engine_contract():
"""The derived flat keys are pure functions of the interval set and
match today's result-key semantics (coder6's mapping)."""
match the engine's result-key semantics."""
ledger = Ledger()
with ledger.span("prior_pass", phase="prior"):
pass
Expand Down
Loading