maybe_defer_rendering

penzai.treescope.foldable_representation.foldable_impl.maybe_defer_rendering(main_thunk: Callable[[RenderableTreePart | None], RenderableTreePart], placeholder_thunk: Callable[[], RenderableTreePart]) → RenderableTreePart

Possibly defers rendering of a part in interactive contexts.

This function can be used by advanced handlers and autovisualizers to delay the rendering of “expensive” leaves such as jax.Array until after the tree structure is drawn. If run in a non-interactive context, this just calls the main thunk. If run in an interactive context, it instead calls the placeholder thunk, and enqueues the placeholder thunk to be called later.

Rendering can be performed in a deferred context by running the handlers under the collecting_deferred_renderings context manager, and then rendered to a sequence of streaming HTML updates using the display_streaming_as_root function.

Note that handlers who call this are responsible for ensuring that the logic in main_thunk is safe to run at a later point in time. In particular, any rendering context managers may have been exited by the time this main thunk is called. As a best practice, handlers should control all of the logic in main_thunk and shouldn’t recursively call the subtree renderer inside it; subtrees should be rendered before calling maybe_defer_rendering.

Parameters:

• main_thunk – A callable producing the main part to render. If not deferred, will be called with None. If deferred, will be called with the placeholder part, which can be inspected to e.g. infer folding state.

• placeholder_thunk – A callable producing a placeholder object, which will be rendered if we are deferring rendering.

Returns:

Either the rendered main part or a wrapped placeholder that will later be replaced with the main part.