DEBUG_MODE = True

def maybe_chexify(fn):
    return chex.chexify(fn) if DEBUG_MODE else fn

@maybe_chexify
@jax.jit
def train_step(...):
    ...

        __
        __