penzai.deprecated.v1.toolshed

Submodules

penzai.deprecated.v1.toolshed.annotate_shapes	Utility to add shape annotations throughout a model.
penzai.deprecated.v1.toolshed.basic_training	Basic training logic for training simple models.
penzai.deprecated.v1.toolshed.check_layers_by_tracing	Utility to check that a layer has been configured correctly by tracing it.
penzai.deprecated.v1.toolshed.interleave_intermediates	Utility to insert intermediate values into a model's PyTree.
penzai.deprecated.v1.toolshed.isolate_submodel	Utility for capturing and analyzing a small subcomputation in a larger model.
penzai.deprecated.v1.toolshed.jit_wrapper	Utilities for JIT compilation of Penzai models.
penzai.deprecated.v1.toolshed.lora	Low-rank adaptation (LoRA, Hu et al. 2021).
penzai.deprecated.v1.toolshed.model_rewiring	Helper classes for rewiring, ablating, and intervening on model activations.
penzai.deprecated.v1.toolshed.sharding_util	Utilities for working with sharded arrays and parameters in Penzai.
penzai.deprecated.v1.toolshed.unflaxify	Recursively transforms flax.linen modules into Penzai layers.