Show HN: AI-related Jax module (I hate if) (github.com)

optimizers.py -> using algebraic expressing without using 'if'

geometry.py -> I leveled it out

math_guardrails.py -> facilitates learning through slight tiliing

jax it truly precious it cures my insomnia

For optimizers.py, are you relying purely on boolean arithmetic for these masks, or did you find a more elegant mathematical workaround to avoid the branching overhead entirely?