Pin user-supplied floats to canonical dtype at every API boundary

[hmgaudecker]

Context

Continues the int-side normalisation in #340 with the float side, though for a completely different reason.

The constraint:

def borrowing_constraint(
    consumption: ContinuousAction,    # action grid, fp32 (quantized via jnp.float32)
    cash_on_hand: FloatND,
    consumption_floor: float,         # Python float — fp64
    equivalence_scale: FloatND,
) -> BoolND:
    return consumption <= cash_on_hand + consumption_floor * equivalence_scale

fired when it should not have -- transfers always ensure that the consumption floor is supported.

What lands on each side without dtype barriers (under jax_enable_x64=True,
which aca_model/__init__.py sets at import):

• LHS consumption: action grid quantized to jnp.float32 in the
  runtime-consumption-points path. Promoted to fp64 for the comparison —
  but promotion preserves the quantization error, it doesn't undo it.
• RHS consumption_floor * equivalence_scale: consumption_floor is a
  Python float (fp64 precision), so the RHS keeps fp64 throughout.

When cash_on_hand took large negative values, the two sides differ by less than
the smallest gap fp64 can represent at that magnitude (a fraction of a single
fp32 quantization step, leaked into fp64 by the promotion). <= flips, and
validate_initial_conditions raises InvalidInitialConditionsError.

This was very annoying to debug. To have one less thing to worry about, this PR makes sure all floats have a consistent dtype.

Overview

Adds canonical_float_dtype() and safe_to_float_dtype next to the int
helpers from #340, and applies them at the same boundaries (params,
initial conditions, transition outputs, V-arrays).

What lands

src/lcm/dtypes.py

• canonical_float_dtype() returns jnp.float64 under
  jax_enable_x64=True, else jnp.float32. Read at call time.
• safe_to_float_dtype(value, *, name) casts to the canonical dtype
  and raises OverflowError (with the leaf's qualified name) when
  down-casting a value above float32 magnitude. Up-casts and
  same-width casts skip the range check; precision loss within range
  is not an error.

Params boundary (src/lcm/params/processing.py)

• _cast_leaves_to_canonical_dtype extends the int-only pass from
  Model.n_subjects: AOT-compile simulate, lock integer dtype to int32 #340 to also cast typed float arrays. Python scalars (int / float
  / bool) pass through to keep JAX weak-typing. pd.Series leaves
  pass through too — convert_series_in_params reshapes them later
  based on the multi-index.

Simulate boundary (src/lcm/simulation/initial_conditions.py)

• build_initial_states casts continuous user arrays to the canonical
  float dtype and pins the missing-state nan fallback to the same
  dtype. Together with the discrete-state int32 cast from Model.n_subjects: AOT-compile simulate, lock integer dtype to int32 #340, the
  simulate state pool has a stable abstract signature across all
  periods.

Transition boundary (src/lcm/simulation/transitions.py)

• _update_states_for_subjects unconditionally casts
  next_state_values to the storage dtype. The cross-kind escape
  hatch added in Model.n_subjects: AOT-compile simulate, lock integer dtype to int32 #340 (so an int-typed user initial condition for a
  continuous state would not be coerced) is no longer needed — the
  initial-state cast above pins storage to the canonical float dtype
  upstream of this site.

Tests

• tests/test_float_dtype_invariants.py (10 tests): helper round-trips,
  initial-state casts, params casts, grid materialisation, V-array
  dtype, multi-period state-dtype stability.
• tests/test_dtypes.py: 7 additional float-helper unit tests.
• tests/test_validate_param_types.py:
  numpy_array_param_rejected -> numpy_array_param_accepted_and_ normalised. With the boundary cast in place numpy arrays are
  auto-converted; the historical rejection-by-isinstance is obsolete.

928 pass, 5 skip; prek + ty clean.

Stacked on

This branch is stacked on feat/simulate-aot-n-subjects (#340) — the
base ref for this PR. Merge order: #340 first, then this. The diff
view here only shows the float-side changes.

Out of scope

• Renaming IntND / FloatND aliases (see Model.n_subjects: AOT-compile simulate, lock integer dtype to int32 #340 review thread —
  permanently out).
• Changing pylcm's default precision posture (conftest still enables
  x64 by default; this normalises within whichever mode the user
  picked).

[mj023]

Looks good, definitely better to use the Jax Types everywhere.

[mj023]

Is there a reason not to use the Age Type anymore?