Mahler-Yum example: split the PRNG key once; drop an unsupported causal claim

Round-3 Pro audit findings, both verified at source.

F2. `create_inputs` drew `type_indices` from `key` and then split that same
already-consumed key for the health, adjustment-cost and productivity-shock
draws. JAX's contract is split-before-consume, and its own key-reuse checker
rejects this exact pattern: the type draw and the other three are not guaranteed
independent. The root key is now split once into four children, each consumed
exactly once. No bias was measured at this calibration (a million-draw
correlation test found none), but independence is what the initial distribution
assumes, so it should hold by construction rather than by luck.

This changes every initial draw, so `tests/test_mahler_yum_2024.py`'s expected
labor-supply distribution and mean-wealth profile at seed=32/n=10000 are now
stale and need regenerating on a GPU. Not done here.

F5. The module docstring claimed the Fortran's old-habit continuation was a
shortcut FORCED by the continuous effort search landing off the habit grid, since
the code could not interpolate along that axis. That is refuted by the authors'
own simulator: `mod_simul.f90` carries `effhabit` as `real(8)`, sets the next
habit to the continuous golden-section effort (line 634), and calls
`bilini(agrid, fgrid, ...)` at the continuous incoming habit throughout. The
two-dimensional interpolation the solver would have needed already existed a few
files away. The mechanism (the effort objective interpolates over assets only)
stands; the motive does not follow from the source and is no longer asserted.

Also softened the numerical-scope note: the moments' numerical derivative not
converging is a property of this implementation, but the measurements do not
isolate WHICH numerical choice causes it. The grid is a consistent explanation,
not an identified one.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01PVSLmeMo6iWUTcNER7bvM6

Author: hmgaudecker

src/lcm_examples/mahler_yum_2024/__init__.py

@@ -19,6 +19,16 @@
 *simulation* meanwhile does advance the habit (`mod_simul.f90:634`), so its
 solve and simulate are mutually inconsistent.
 
+What their solver omits is an interpolation along the habit axis: with
+`puregrid = 2` the chosen effort lands off the 30-point habit grid (`nf = 30`),
+so evaluating `EV` there needs one, and inside the effort objective they
+interpolate over assets only (`lini(agrid, EV(:,h), ap)`). That is the mechanism;
+the source does not reveal the motive, and it does not support calling the
+omission *forced*. Their own simulator carries `effhabit` as `real(8)`, sets the
+next habit to the continuous golden-section effort (`mod_simul.f90:634`), and
+calls `bilini(agrid, fgrid, ...)` at that continuous habit throughout — the
+two-dimensional interpolation the solver would have needed already existed.
+
 `next_lagged_effort` below implements eq. (6). Consequently this model must NOT
 be expected to reproduce the paper's published model column at the published
 Table-II estimates: those estimates were fit against the other recursion.
@@ -44,10 +54,12 @@
 5-node equal-probability rule integrates it exactly only when a single branch
 dominates over the whole support; the node count needs a refinement check. Note
 that simulated adjustment shares are NOT restricted to multiples of 1/5 —
-simulation draws the cost continuously. (ii) The discrete grids make the
-simulated moments locally flat in the parameters, which is why MSM inference is
-withheld in the replication. That is a property of this solver, not of the
-paper's model: the authors' continuous searches do not share it.
+simulation draws the cost continuously. (ii) The simulated moments' numerical
+derivative does not converge across finite-difference steps, which is why MSM
+inference is withheld in the replication. The instability is a property of this
+implementation, not of the paper's model — the authors searched continuously and
+estimated successfully — but the measurements do not isolate *which* numerical
+choice produces it; the grid is a consistent explanation, not an identified one.
 """
 
 import dataclasses
@@ -933,18 +945,25 @@ def create_inputs(
     }
 
     td = _build_type_distribution()
-    key = jax.random.key(seed)
+    # Split the root key ONCE and consume each child exactly once. Deriving
+    # `type_indices` from `key` and then splitting that same `key` reuses a
+    # consumed key, which JAX's key-reuse checker rejects: the type draw and the
+    # three draws below are then not guaranteed independent. No bias was measured
+    # at this calibration, but the guarantee is what the initial distribution
+    # relies on, so it has to hold by construction rather than by luck.
+    type_key, health_key, adjustment_key, shock_key = jax.random.split(
+        jax.random.key(seed), num=4
+    )
     type_indices = np.asarray(
         jax.random.choice(
-            key,
+            type_key,
             jnp.arange(len(td)),
             (n_simulation_subjects,),
             p=jnp.array(td["probability"].to_numpy()),
         )
     )
 
-    keys = jax.random.split(key=key, num=3)
-    health_draw = jax.random.uniform(keys[0], (n_simulation_subjects,))
+    health_draw = jax.random.uniform(health_key, (n_simulation_subjects,))
     health_thresholds = np.asarray(td["health_threshold"])[type_indices]
 
     # Round the continuous initial habit to the NEAREST grid point. `searchsorted`
@@ -978,15 +997,15 @@ def create_inputs(
             "productivity_shock": np.asarray(
                 shock_gridpoints[
                     jax.random.choice(
-                        keys[2],
+                        shock_key,
                         jnp.arange(5),
                         (n_simulation_subjects,),
                         p=stationary_shock_dist,
                     )
                 ]
             ),
             "adjustment_cost": np.asarray(
-                jax.random.uniform(keys[1], (n_simulation_subjects,))
+                jax.random.uniform(adjustment_key, (n_simulation_subjects,))
             ),
         }
     )