run_and_interleave_intermediates

penzai.toolshed.interleave_intermediates.run_and_interleave_intermediates(root: pz.Layer, argument: Any, jit: bool = True) → pz.Layer

Interleaves intermediate values into a model.

Parameters:

• root – The layer or model to annotate.

• argument – An argument to run the model with.

• jit – Whether to run the model under jax.jit.

Returns:

A copy of the layer or model with saved intermediates added before and after every layer that was called during the evaluation of the model. Saved intermediates are represented as IdentityWithSavedActivations layers, which hold onto saved activations but do not use them.