tuned_lens.plotting.trajectory_plotting#

Contains utility classes for creating heatmap visualizations.

Functions

# empty test needed in case the module has no example usage.
# otherwise, testsetup throws an error
pass

tuned_lens.plotting.trajectory_plotting.trunc_string_left(string, new_len)#

Truncate a string to the left.

Return type:: str

Classes

class tuned_lens.plotting.trajectory_plotting.TrajectoryLabels(label_strings, hover_over_entries=None)#

Contains sets of labels for each layer and position in the residual stream.

hover_over_entries: Optional[ndarray[Any, dtype[str_]]] = None#: (n_layers x sequence_length x rows x cols) table of strings to display when hovering over a cell. For example, the top k prediction from the lens.

label_strings: ndarray[Any, dtype[str_]]#: (n_layers x sequence_length) label for each layer and position in the stream.

stride(stride)#

Return a new TrajectoryLabels with the given stride.

Parameters:: stride – The number of layers between each layer we keep.
Return type:: TrajectoryLabels
Returns:: A new TrajectoryLabels with the given stride.

template_and_customdata(col_width_limit=10)#

Construct a template for use with Plotly’s hovertemplate.

Return type:: Tuple[str, ndarray[Any, dtype[str_]]]

class tuned_lens.plotting.trajectory_plotting.TrajectoryStatistic(name, stats, sequence_labels=None, trajectory_labels=None, units=None, max=None, min=None, includes_output=True, _layer_labels=None)#

This class represents a trajectory statistic that can be visualized.

For example, the entropy of the lens predictions at each layer.

clip(min, max)#

Return a new TrajectoryStatistic with the given min and max.

Parameters:

min – The minimum value to clip to.
max – The maximum value to clip to.

Return type:

TrajectoryStatistic

Returns:

A new TrajectoryStatistic with the given min and max.

figure(title='', colorscale='rdbu_r', token_width=80)#

Produce a heatmap plot of the statistic.

Parameters:

title – The title of the plot.
colorscale – The colorscale to use for the heatmap.
token_width – The width of each token in the plot.

Return type:

Figure

Returns:

The plotly heatmap figure.

heatmap(colorscale='rdbu_r', log_scale=False, **kwargs)#

Returns a Plotly Heatmap object for this statistic.

Parameters:

colorscale – The colorscale to use for the heatmap.
log_scale – Whether to use a log scale for the colorbar.
**kwargs – Additional keyword arguments to pass to the Heatmap constructor.

Return type:

Heatmap

Returns:

A plotly Heatmap where the x-axis is the sequence dimension, the y-axis is the layer dimension, and the color of each cell is the value of the statistic.

includes_output: bool = True#: Whether the statistic includes the final output layer.

max: Optional[float] = None#: The maximum value of the statistic.

min: Optional[float] = None#: The minimum value of the statistic.

name: str#: The name of the statistic. For example, “entropy”.

sequence_labels: Optional[ndarray[Any, dtype[str_]]] = None#: (sequence_length) labels for the sequence dimension e.g. input tokens.

stats: ndarray[Any, dtype[float32]]#: (n_layers x sequence_length) value of the statistic across layer and position.

stride(stride)#

Return a new TrajectoryStatistic with the given stride.

Parameters:: stride – The number of layers between each layer we keep.
Return type:: TrajectoryStatistic
Returns:: A new TrajectoryStatistic with the given stride.

trajectory_labels: Optional[TrajectoryLabels] = None#: Labels for each layer and position in the stream. For example, the top 1 prediction from the lens at each layer.

units: Optional[str] = None#: The units of the statistic.