What is egregore-core-jax?

egregore-core-jax is a small collection of Python utility modules for people building deep learning models with JAX, a numerical computing library that compiles code to run efficiently on GPUs and TPUs. The project focuses on a specific kind of performance problem: ordinary conditional logic, if this then that, can slow down JAX programs because it interrupts the compiler's ability to fuse operations into fast, uninterrupted chunks of hardware instructions. Each module in this repository replaces that kind of conditional logic with pure arithmetic instead. The first module, optimizers.py, helps apply different learning rates to different layers of a model, a common technique called layer-wise learning rate decay. Rather than checking each parameter's name against a list of rules with if statements, it turns those checks into numeric masks that get multiplied against the parameters directly, which keeps the whole operation compiling into a single fast unit instead of many small conditional steps. The second module, geometry.py, deals with a common annoyance in JAX: code that reshapes tensors of varying dimensions can trigger confusing compiler errors when the shape changes unexpectedly during tracing. This module locks certain shape related arguments as fixed constants at compile time, so tensors of different shapes, such as a plain batch of vectors or a batch of sequences, can be reshaped and reduced using the same underlying fast code path without triggering those errors. The third module, math_guardrails.py, protects against a numerical stability problem: some math functions, such as arccos, produce invalid results or break gradient calculations when their input gets too close to certain boundary values. This module keeps inputs within a safe range while still allowing a small gradient signal to pass through at the edges, so that training does not stall out completely. Overall this is a low level performance toolkit aimed at engineers already comfortable with JAX internals, gradient computation, and hardware compiler behavior, rather than a general purpose library for beginners. The code is released under the Apache License 2.0, which allows free use, modification, and commercial redistribution.