Compile the SMC trace function once and pass compile_kwargs to it - #8450

Description

After sampling, sample_smc built each chain's trace with a fresh NDArray, so the trace function was compiled once per chain. It was also compiled without the compile_kwargs given to sample_smc, so mode="FAST_COMPILE" only applied to the kernel.

This follows the approach described in the issue:

• The first chain's trace is compiled with compile_kwargs. Later chains reuse its fn, var_shapes and var_dtypes. This covers both the parallel and the sequential path.
• BaseTrace takes an optional compile_kwargs that is forwarded to model.compile_fn. The trace's own settings (point_fn, on_unused_input, trust_input) still take precedence. The SMC kernel adds on_unused_input to the same dict, so they would otherwise clash.

On a 20-variable model with 4 chains, model.compile_fn for the trace now runs once instead of four times.

Related Issue

• Closes BUG: Long post-sampling compile times on large models with SMC sampler #8347
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass (ruff 0.11.13 check and format on the changed files)
• Included tests that prove the fix is effective or that the new feature works: test_smc_compiles_trace_fn_once for cores=1 and cores=2 fails on main (3 compiles, mode=None) and passes with the fix
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

tests/smc and tests/backends pass locally. The exception is test_mv_missing_data_model, which fails on a missing BLAS warning in my environment.

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):