penzai.toolshed

Submodules

penzai.toolshed.auto_nmap	Utility to automatically apply nmap to functions in a module.
penzai.toolshed.basic_training	Basic training logic for training simple models.
penzai.toolshed.gradient_checkpointing	Utilities for gradient checkpointing / rematerialization.
penzai.toolshed.isolate_submodel	Utility for capturing and analyzing a small subcomputation in a larger model.
penzai.toolshed.jit_wrapper	Utilities for JIT compilation of Penzai models.
penzai.toolshed.lora	Low-rank adaptation (LoRA, Hu et al. 2021).
penzai.toolshed.model_rewiring	Helper classes for rewiring, ablating, and intervening on model activations.
penzai.toolshed.patch_ipdb	Patches the ipdb debugger to enable rich output (e.g. with treescope).
penzai.toolshed.save_intermediates	Utilities to capture and hold intermediate values as a model runs.
penzai.toolshed.sharding_util	Utilities for working with sharded arrays and parameters in Penzai.
penzai.toolshed.token_visualization	Helpers to visualize tokens, token probabilities, and token sequences.
penzai.toolshed.unflaxify	Recursively transforms flax.linen modules into Penzai layers.