jit_wrapper

Utilities for JIT compilation of Penzai models.

This module provides a JIT-compilation wrapper that preserves the structure of the original model, and also supports mutable variables inside the jitted block.

The intended use of these wrappers is to enable interactive exploration of models in e.g. Colab notebooks, while still taking advantage of JIT compilation.

Classes

Jitted

Wraps a pure layer to run under JIT compliation.