[cymbalrush]

Add CoreAIProbedInspector: read intermediates as extra model outputs

Why

CoreAIInspector reads intermediates through the runtime's intermediate logger, which does not fire for work placed on the neural engine, so on the ANE it returns None for every op. The ANE does not natively expose intermediates, so the only route today is to make them outputs of the model.

How it works

CoreAIProbedInspector(program, function_name="main", specialization_options=None, max_cached_models=4) implements the usual Inspector interface, so it drops in where CoreAIInspector is used. For get_intermediates_for_ops(op_ids, inputs, state=None):

1. Copy and rewrite. It deep-copies the program and appends each requested op's tensor results to the function's coreai.output. It updates the function type and result names, so the model's own outputs keep their names and positions.
2. Compile and run. It saves, compiles and runs the copy with the options you chose, for example a preferred compute unit.
3. Map back. Results come back keyed by Core AI op ID, one array per result. Ops that cannot be probed return None and are logged.

The rewrite lives in probes.py and does not need the Python runtime, so it can be reused elsewhere:

• add_probe_outputs(program, op_ids, function_name) returns a ProbedProgram.
• split_outputs(arrays) separates a run's outputs into the model's own and the probed values.
• manifest() returns the probe names as JSON-able values.
• save(dir) writes probed.aimodel and probes.json for another runtime.

Other behaviors:

• State. Stateful models work. Pass state= by name. The inspector copies it on each call, so every call starts from the same values and your arrays are left alone. last_state holds what the run left behind.
• Compile cache. Compiled copies are kept per set of probed op IDs, least recently used dropped first. max_cached_models sets the limit (default 4). 0 compiles on every call, and clear_cache() drops everything.
• last_outputs. It returns only the model's own outputs, without the probe outputs.

Be careful: probed values are not the original model's values

Every probe changes the model, so treat results as error-prone.

• A probe is an output. The compiler cannot fuse it away, so it may split or re-place work. Probed values can differ from what the unprobed model computes, and the ANE may handle the extra outputs differently.
• ANE values are fp16-ANE values. In a small test they were within about 3e-4 relative error for conv and relu, and about 2.5% for sigmoid×x and the conv after it. Compare against a reference of the same precision, or use a loose tolerance.
• Check each probe set before relying on it. Add a small set, then compare the final output and compute-plan placement (plan_diff) against an unprobed run. If either changes, those values cannot be trusted. Reduce the set or try another. If they match, they are a reasonable stand-in. This is a manual check; the inspector does not run it.
• Prefer fewer, coarser probes. Probing every op perturbs the model most and costs the most compiles.
• Cached models are tied to the exact probe set. A copy compiled for a larger set is not reused to answer a smaller one.
• Each new probe set costs a compile. A big model is expensive.
• Limits:
  • Only tensor results of ops in the function's top-level graph can be probed. Ops inside invoked graphs, ops without tensor results, and unknown IDs return None.
  • The program must not be mutated after the inspector is created, or the cache goes stale.

Not included

Module-level probing, which would choose probe points at module boundaries, is not part of this PR. Module attribution is incomplete for some ops. This PR probes whichever op IDs it is given.

Testing

tests/debugging/test_probed_inspector.py (7 tests) covers:

• the original program and its outputs are left alone;
• unprobeable IDs are reported;
• intermediates are returned;
• stateful models are read without advancing the state;
• the manifest, split and save work without a runtime;
• cache reuse, the size limit, and switching the cache off.

They pass on the default runtime. I also ran the first 5 on the neural-engine setting, but the cache tests were added afterwards and have only run on the default runtime. They use a small fp16 toy model. I have not run this on a large model or a real KV-cache model, and I have not measured the compile-time saving.