Cache jit closures instead of re-tracing on every call

[Codcore]

Closes #136.

__jit__/5 forwarded straight to __compile__/4, which discards the key Nx
supplies and calls native_compile/3. That runs fun.(vars) and walks the
resulting expression on every call — only to arrive at a program already sitting
in the program table.

This caches the closure __compile__/4 returns, keyed by that key plus the
shapes and types of vars and the device. On a hit, the trace never runs.

What is left alone

The program table and its structural key. That key is what lets structurally
identical layers share one compiled program across different call sites, and it
still does. The new table sits above it and short-circuits repeat calls of the
same call site with the same argument shapes — the shape of every decode loop and
every training step.

Calls carrying :command_queue or :hooks are not cached: those options
reference caller-owned state, and a process-wide table should not extend its
lifetime. Such calls take the previous path unchanged.

Measurements

A chain of depth cheap ops, so the timing reflects graph traversal rather than
arithmetic (script in #136):

  depth   before    after    Nx.Defn.compile
     32   0.59 ms   0.43 ms   0.37 ms
    128   1.01 ms   0.53 ms   0.38 ms
    512   3.12 ms   1.18 ms   0.67 ms
   1024   5.89 ms   1.85 ms   1.30 ms

At depth 1024 that is 3.2× faster, and what remains between jit and compile
is 0.55 ms rather than 4.48 ms.

On a real workload — a 40-layer Qwen3.6-35B-A3B MoE decode step, roughly 4300
nodes — the change removes 8.35 ms per token out of 89.

Tests

mix test: 2668 passed (825 doctests, 1843 tests), 7 excluded. No failures.

A note on the trade-off

The cache is unbounded, like the program table it sits above. Each entry holds a
closure referencing a compiled program, so a caller that jits the same function
across many distinct argument shapes will accumulate entries. If that is a
concern, the same key would work with a bounded table; I kept it simple to match
the existing behaviour rather than introduce a second eviction policy.

[polvalente]

I think hooks is always a given option. Not sure why this is being filtered on.

[polvalente]

You also want to take into account the vectorized_axes and the donation arguments. using Nx.Defn.Composite.traverse with Nx.to_template is likely a better conversion for the arguments. I think the options should also be taken into account because hooks might differ between calls