penzai.toolshed

Submodules

penzai.toolshed.annotate_shapes	Utility to add shape annotations throughout a model.
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.check_layers_by_tracing	Utility to check that a layer has been configured correctly by tracing it.
penzai.toolshed.interleave_intermediates	Utility to insert intermediate values into a model's PyTree.
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.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.