Pin jax>=0.4.1 to fix scipy.linalg.tril AttributeError

JAX 0.3.0 (resolved for py310) references scipy.linalg.tril which was
removed in newer scipy. Pin to >=0.4.1 to ensure compatibility.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: timmens