Expose prefill/decode context caching via fit_mode="fit_with_cache" (JAX)

[jjovalle99]

Refs: #60

What

Adds an opt-in fit_mode="fit_with_cache" to TabFMClassifier and
TabFMRegressor (JAX backend):

• fit() runs model.prefill once per ensemble member view, chunked by
  batch_size, and stores the context caches on the estimator.
• predict and predict_proba run model.decode for the query rows against
  the cached context, instead of re-encoding context plus queries on every
  call.
• The default fit_mode="fit" is unchanged. The cached path reuses the same
  class-shift correction and logit post-processing as the default path.

Why

Predict currently re-processes the full training context on every call, once
per ensemble member. prefill/decode already exist in the JAX model
(covered by model_test.py::test_prefill_decode_consistency) and were built
to move that work to fit time. This PR exposes them through the sklearn API.

Steady-state measurements on the released 1.6B classification checkpoint
(RTX 5090 32 GB, bf16, creditcard dataset from OpenML, n_estimators=1,
10k-row context; the one-time prefill inside fit takes 28.9 s):

queries per predict	fit_mode="fit"	fit_mode="fit_with_cache"	speedup
1	2.51 s	1.93 s	1.3x
100	2.51 s	2.10 s	1.2x
1,000	2.94 s	2.13 s	1.4x
10,000	10.24 s	4.25 s	2.4x

The default path's per-call cost grows faster than linearly with context
length: per member, about 2.5 s at 10k rows, 9 s at 20k, 35 s at 40k. The
cached path's cost does not depend on context length, so the benefit grows
exactly where TabFM is slowest. The table stops at 10k because larger cached
contexts hit decode's memory limits (see Known limitations).

Equivalence

• fp32 unit tests: max logit difference about 7e-7 and max probability
  difference about 8e-8 between the cached and default paths, including the
  ensemble() preset (NNLS plus calibration), categorical permutation, and
  row subsampling.
• bf16, released checkpoint, 10k test rows: mean probability difference 1e-4
  to 2e-4. Max difference 0.012 with n_estimators=4, and 0.134 on isolated
  rows with n_estimators=1, where no ensemble averaging damps bf16 noise.
  Both sit within the atol=0.2 that test_prefill_decode_consistency uses
  for bf16-scale numerics. The paths are numerically equivalent, not bitwise
  identical.

Implementation notes

• decode runs eagerly, matching how model_test.py uses it. Wrapping it in
  nnx.jit copies the multi-GB caches into the executable arena (about 2x
  peak memory) and fuses a transpose across all layers that runs out of
  memory or fails XLA autotuning at real context lengths. The caches returned
  by prefill are passed to decode unmodified, so the wrapper never touches
  ICLearningCache internals.
• A new EnsembleGenerator.context_tensors() builds the fitted context-only
  member tensors directly from generator state. The public transform path
  and its sklearn validation are untouched. A unit test checks its output,
  element for element, against the transform-and-slice derivation under
  permute_categorical=True and max_num_rows.
• Caches scale with context length times members. Members are chunked by
  batch_size at prefill and decode, and caches are stored per chunk.
• Caches are dropped on pickle and rebuilt lazily on the next predict,
  following the existing _COMPILED_PREDICT_CACHE_ATTRS pattern.
• CHANGELOG.md is left untouched; it appears to be maintainer-managed
  through the Copybara sync.

Known limitations

• The cached decode path runs on a single device and raises
  NotImplementedError under a multi-device data mesh. The column cache's
  leading axis is the layer axis, not batch, so the default sharding spec
  does not apply. Sharding support is a follow-up.
• decode's memory use bounds the usable context on consumer GPUs. On a 32 GB
  card, cached predict works at a 10k context with one member, and runs out
  of memory at 10k with four members and at 20k or more. The cause is
  decode's internal per-head key/value tensor spanning all 24 layers
  ([heads, layers, B, T, head_dim], about 1 GB per 10k context rows per
  member); at 40k the same tensor also fails XLA autotuning. This is a
  property of the existing decode, not of this change: the same shapes
  appear when calling decode directly. Chunking that tensor per layer would
  unlock larger contexts. We can file this separately.
• The PyTorch backend has no prefill/decode yet, so this is JAX only.

Tests

• New CachedPredictTest (13 tests): probability and label parity against
  the default mode, including categorical permutation, row subsampling,
  single-chunk, and ensemble() preset configs; regressor value parity;
  pickle round-trips for both estimators; refit rebuilds the cache; a
  positive check that the decode path actually runs; invalid fit_mode; and
  the non-JAX and multi-device guards.
• New EnsembleGeneratorTest::test_context_tensors_matches_dummy_query_derivation.
• Full wrapper suite: 60 passed (46 baseline plus 14 new), no regressions.
  model_test.py untouched and passing.