Yeah, Jax with Equinox, jaxtyping, and leaning hard on python’s static typing modules + typeguard lets you pretend that you have a nice little language embedded in python. I swore off Julia a few years ago.