penzai.experimental.v2.toolshed

Submodules

penzai.experimental.v2.toolshed.basic_training	Basic training logic for training simple models.
penzai.experimental.v2.toolshed.gradient_checkpointing	Utilities for gradient checkpointing / rematerialization.
penzai.experimental.v2.toolshed.isolate_submodel	Utility for capturing and analyzing a small subcomputation in a larger model.
penzai.experimental.v2.toolshed.jit_wrapper	Utilities for JIT compilation of Penzai models.
penzai.experimental.v2.toolshed.lora	Low-rank adaptation (LoRA, Hu et al. 2021).
penzai.experimental.v2.toolshed.model_rewiring	Helper classes for rewiring, ablating, and intervening on model activations.
penzai.experimental.v2.toolshed.save_intermediates	Utilities to capture and hold intermediate values as a model runs.
penzai.experimental.v2.toolshed.sharding_util	Utilities for working with sharded arrays and parameters in Penzai.
penzai.experimental.v2.toolshed.unflaxify	Recursively transforms flax.linen modules into Penzai layers.