ParameterCount cached at construction, stale after lazy shape resolution

[ooples]

Problem

NeuralNetworkBase<T>.ParameterCount is computed at construction time and never invalidated after lazy input-shape resolution. Any model whose layers use lazy-shape sentinels (in=[1] at construction, resolved to real shapes on first forward via positional encoding, skip-concat, etc.) hits an inconsistency where ParameterCount reports the pre-resolution value while GetParameters() walks layers fresh and returns the post-resolution value.

The concrete failure mode: SetParameters(loaded) reads the stale ParameterCount for length validation, rejects a correctly-sized saved vector, throws ArgumentException: Expected N parameters, got M.

This was originally filed as "hierarchical sampling breaks save/load" — the real bug is broader.

Diagnostic data (AiDotNet 0.224.3-preview1826.3 / #1829 merged into master)

Repro: construct NeRF<float>(hiddenDim: 192, numLayers: 8, positionEncodingLevels: 10, directionEncodingLevels: 4), snapshot before + after one Train(dummyInput, dummyTarget) call.

State	Layers.Count	ParameterCount	GetParameters().Length	GetParameterChunks sum
Fresh (hierarchical=false)	12	285,444	285,444	285,444
After 1 Train (hierarchical=false)	12	285,444	348,036	348,036
Fresh (hierarchical=true)	12	285,444	285,444	285,444
After 1 Train (hierarchical=true)	12	285,444	348,036	348,036

Delta = 62,592 params, identical with hierarchical ON and OFF, so this is not hierarchical-specific.

Per-layer input shapes that resolve after the first Train:

Layer index	Before	After	Δ params	Cause
[0] input	in=[1] (384)	in=[60] (11,712)	+11,328	Positional encoding of position: 10 levels × 3 dims × 2 (sin+cos) = 60
[4] skip	in=[192] (37,056)	in=[252] (48,576)	+11,520	Skip-connection concat with 60-dim encoded position
[9] color input	in=[1] (384)	in=[192] (37,056)	+36,672	Feature vector fed from density head
[10] direction	in=[192] (24,704)	in=[216] (27,776)	+3,072	Positional encoding of direction: 4 levels × 3 dims × 2 = 24

ParameterCount (285,444) is stale — reflects layer weight allocations sized against the [1] sentinels, not the resolved shapes. GetParameters() and GetParameterChunks() correctly walk layers fresh and return 348,036.

Impact

• Blocks save/load for any radiance-field model (NeRF regardless of hierarchical, likely InstantNGP and GaussianSplatting similarly). Downstream: Week 11 Session B NerfLab demo cannot pre-bake trained weights; forced to retrain live every run.
• Latent for other consumers that trust ParameterCount as an authoritative size — HPO tracking, checkpointing, mixed-precision buffer sizing, distributed parameter sharding.

Fix direction (agreed in design discussion)

Three coordinated pieces — the combination exceeds industry standard rather than merely matching it:

B — Version-invalidated ParameterCount cache

Keep the cached property for hot-path access (optimizer step reads it every batch), but bump a monotonic version whenever Layers is mutated OR any layer's shape resolves. First read after a mutation recomputes; subsequent reads hit cache. Correctness on par with PyTorch's re-walk-every-call; performance strictly better on repeated access.

Wiring: version-bump hooks in Layers.Add / Insert / Remove, in Layer.SetInputShape (and equivalent lazy-resolution sites), and any code that mutates the trainable set. ConfigureOptimizer is the natural extension surface for exposing the invalidation policy — nullable option, default policy is "auto-invalidate on any structural mutation" (industry-standard behavior), advanced callers can override.

D — Shape-aware save/load

Saved payload includes per-layer shape metadata (input/output shape, ParameterCount) alongside the flat weight vector. SetParameters matches by named-layer + shape, not by flat total. If a target layer has a lazy sentinel, load auto-materializes it to the saved shape. This eliminates the mismatch failure mode entirely — not by making the count match, but by making the check smarter.

Wiring: extend ConfigureCheckpointing. Nullable options in a checkpoint-config object:

• SchemaVersion — checkpoint format version tag (industry-standard default: 1)
• IncludePerLayerShapes — default true (industry-standard: on)
• IncludeOptimizerState — default true (industry-standard: on; the combined-artifact win)
• AutoMaterializeOnLoad — default true (industry-standard: on)
• Custom serializers / migrators for backward-compat on schema-version bumps

E — Explicit ResolveShapes(sampleInput) opt-in

Add NeuralNetworkBase<T>.ResolveShapes(Tensor<T> sampleInput) — a public method that runs a dry-run forward, resolves all lazy shapes, bumps the version cache, and locks the structure. Callers who want up-front determinism (Keras-style .compile(), but as a first-class opt-in) invoke it before the first Train; callers who prefer lazy semantics ignore it entirely.

Wiring: ConfigureModel(model, sampleInput: ...) gets an optional sampleInput parameter. When supplied, the facade calls model.ResolveShapes(sampleInput) immediately after ConfigureModel — so BuildAsync sees a fully-materialized model and the caller can inspect GetParameters().Length before training. Neither PyTorch nor Keras exposes this as an opt-in method — this is the beyond-industry step.

Why B + D + E together exceed industry standard

• PyTorch walks the tree on every parameters() call — correct but not cached. B beats it on hot-path.
• PyTorch uses named-key state_dict / load_state_dict — good but doesn't auto-materialize lazy shapes on load. D beats it.
• Keras exposes model.compile() for up-front shape resolution — but that's tied to the training-graph compile step, not offered as a standalone method. E offers the same semantics decoupled from training.
• No mainstream framework offers a combined save-load artifact (model + optimizer + step) as the industry-standard default. D does.

Related

• Facade BuildAsync throws on radiance-field models: feature-selection guard misses 1-D-encoded neural inputs #1826 / fix(facade): unblock BuildAsync for radiance-field models (#1826) #1829 — original facade unblock work that surfaced this while attempting a save/load round-trip
• Follow-ups likely needed for the same lazy-shape pattern in InstantNGP and GaussianSplatting once the base-class fix lands