SOMOrderer.order cannot be put under jax.jit or jax.vmap. grad works, and the expensive inner step is already jitted, so this is a capability gap rather than a performance one — but it is the only orderer with the limitation, and it makes a chain containing it unjittable too.
Measured
With metric_scale actually traced (a test function that ignores its argument constant-folds and misleadingly passes):
|
jit |
grad |
LocalFlowOrderer alone |
✅ |
✅ |
LocalFlowOrderer | LocalFlowOrderer |
✅ |
✅ |
SOMOrderer alone |
❌ TracerBoolConversionError |
✅ |
LocalFlowOrderer | SOMOrderer |
❌ NonConcreteBooleanIndexError |
✅ |
It fails alone, so this is not a chaining problem — #70 made a chain no less jittable than its stages, and this is the stage.
Where
SOMOrderer.order is eager orchestration around a jitted core. Three constructs concretize:
som.py:321 — work = prior[prior >= 0]. Boolean mask on init.indices, producing a variable-length array. This is the NonConcreteBooleanIndexError on the chained path. (chord_along_ordering in base.py had the same shape and was rewritten with argmax + a scratch-slot scatter — the same trick applies here only if the working-set size can be made static.)som.py:239 — if self.metric_scale: in _resolve_metric. Truth-test on a traced scalar; this is the TracerBoolConversionError when metric_scale is traced.som.py:277 — rho = float(...) then if rho >= _DISAGREE_WARN in _warn_if_disagrees. A host sync plus a Python branch, and warnings.warn cannot fire from inside a trace at all.
som.py:312 (n_obs = int(...)) is fine — it reads a static shape.
The hard part
This is not a mechanical fix. The working set is data-dependent: work = prior[prior >= 0] has a length that depends on how many observations the previous stage visited, and every downstream shape (sub_q, the prototypes, lam) follows from it. Under jit all of those must be static.
Options, roughly in increasing order of honesty:
- Fix only what is cheap. Make
_resolve_metric branch on a static flag rather than the scalar's truthiness, and make _warn_if_disagrees a no-op under trace. That clears the TracerBoolConversionError but not the masking, so jit still fails when chained.
- Pad instead of compact. Keep the full-length arrays and carry a validity mask through
fit/chord, weighting unvisited observations to zero. Shapes stay static and the whole thing traces. Costs work proportional to n_obs rather than n_visited, and touches the SOM core.
- Declare it eager. Document that
SOMOrderer.order is an eager driver whose heavy step (_train_and_project) is jitted internally, and that user-level jit is not supported. Cheapest, and arguably honest — but it makes the orderer family inconsistent.
Worth asking first
Does anything actually need it? The cost is in fit/densify/chord, and _train_and_project already wraps all three in eqx.filter_jit. User-level jit would mainly buy composition — putting an ordering inside a larger traced computation — and vmap over hyperparameters. If neither is wanted, (3) is the right answer and this issue should close as documentation.
The ensemble work in #58 would want it: a Statistically Combined Ensemble of SOMs is a vmap over members, and phasecurvefit.som's functional core was deliberately written to support that. Note the core is already traceable — it is the SOMOrderer wrapper that is not — so #58 may be servable without this.
Follows from #56, #70.
SOMOrderer.ordercannot be put underjax.jitorjax.vmap.gradworks, and the expensive inner step is already jitted, so this is a capability gap rather than a performance one — but it is the only orderer with the limitation, and it makes a chain containing it unjittable too.Measured
With
metric_scaleactually traced (a test function that ignores its argument constant-folds and misleadingly passes):jitgradLocalFlowOrdereraloneLocalFlowOrderer | LocalFlowOrdererSOMOrdereraloneTracerBoolConversionErrorLocalFlowOrderer | SOMOrdererNonConcreteBooleanIndexErrorIt fails alone, so this is not a chaining problem — #70 made a chain no less jittable than its stages, and this is the stage.
Where
SOMOrderer.orderis eager orchestration around a jitted core. Three constructs concretize:som.py:321—work = prior[prior >= 0]. Boolean mask oninit.indices, producing a variable-length array. This is theNonConcreteBooleanIndexErroron the chained path. (chord_along_orderinginbase.pyhad the same shape and was rewritten withargmax+ a scratch-slot scatter — the same trick applies here only if the working-set size can be made static.)som.py:239—if self.metric_scale:in_resolve_metric. Truth-test on a traced scalar; this is theTracerBoolConversionErrorwhenmetric_scaleis traced.som.py:277—rho = float(...)thenif rho >= _DISAGREE_WARNin_warn_if_disagrees. A host sync plus a Python branch, andwarnings.warncannot fire from inside a trace at all.som.py:312(n_obs = int(...)) is fine — it reads a static shape.The hard part
This is not a mechanical fix. The working set is data-dependent:
work = prior[prior >= 0]has a length that depends on how many observations the previous stage visited, and every downstream shape (sub_q, the prototypes,lam) follows from it. Underjitall of those must be static.Options, roughly in increasing order of honesty:
_resolve_metricbranch on a static flag rather than the scalar's truthiness, and make_warn_if_disagreesa no-op under trace. That clears theTracerBoolConversionErrorbut not the masking, sojitstill fails when chained.fit/chord, weighting unvisited observations to zero. Shapes stay static and the whole thing traces. Costs work proportional ton_obsrather thann_visited, and touches the SOM core.SOMOrderer.orderis an eager driver whose heavy step (_train_and_project) is jitted internally, and that user-leveljitis not supported. Cheapest, and arguably honest — but it makes the orderer family inconsistent.Worth asking first
Does anything actually need it? The cost is in
fit/densify/chord, and_train_and_projectalready wraps all three ineqx.filter_jit. User-leveljitwould mainly buy composition — putting an ordering inside a larger traced computation — andvmapover hyperparameters. If neither is wanted, (3) is the right answer and this issue should close as documentation.The ensemble work in #58 would want it: a Statistically Combined Ensemble of SOMs is a
vmapover members, andphasecurvefit.som's functional core was deliberately written to support that. Note the core is already traceable — it is theSOMOrdererwrapper that is not — so #58 may be servable without this.Follows from #56, #70.