https://hamon.dev/tags/jax/ Recent content in Jax on so i built a thing Hugo en Thu, 04 Jun 2026 00:00:00 +0000 https://hamon.dev/blog/lumpyspace4/ Thu, 04 Jun 2026 00:00:00 +0000 https://hamon.dev/blog/lumpyspace4/ <h2 id="a-confession"> A Confession <a class="heading-link" href="#a-confession"> <i class="fa-solid fa-link" aria-hidden="true" title="Link to heading"></i> <span class="sr-only">Link to heading</span> </a> </h2> <p>In my <a href="https://hamon.dev/blog/lumpyspace3/" >previous post</a>, I explained that we had implemented an Augmented Lagrangian (AL) method to strictly enforce the Weak Energy Condition (WEC). I shared monitoring showing the WEC loss dropping to exactly <code>0.000000e+00</code>.</p> <p>However, it turns out we had fallen victim to a silent scoping bug inside the JAX <code>@eqx.filter_jit</code> compiler.</p> <p>To capture static hyperparameters (like <code>w_efe</code>, <code>w_sn</code>) without passing them explicitly, we wrote our <code>loss_fn</code> as an internal closure inside the compiled <code>step</code> function. This is a common JAX pattern. However, the Augmented Lagrangian multiplier ($\lambda_{\text{WEC}}$) is a dynamic array that changes every single optimization step.</p>