devinterp package

Subpackages

Submodules

devinterp.utils module

class devinterp.utils.Outputs(loss)

Bases: tuple

count(value, /)

Return number of occurrences of value.

index(value, start=0, stop=9223372036854775807, /)

Return first index of value.

Raises ValueError if the value is not present.

loss: Tensor

Alias for field number 0

devinterp.utils.prepare_input(data: Tensor | Any, device, is_deepspeed_enabled=False, accelerator=None) → Tensor | Any

Prepares one data before feeding it to the model, be it a tensor or a nested list/dictionary of tensors.

Adapted from HuggingFace’s transformers’s Trainer._prepare_input().

devinterp.utils.set_seed(seed: int, device: str | device | None = None)

Sets the seed for the Learner.

Parameters:

seed (int) – Seed to set.

devinterp.vis_utils module

class devinterp.vis_utils.EpsilonBetaAnalyzer

Bases: object

A class for analyzing and visualizing the local learning coefficient (LLC) across different epsilon and beta values.

Includes methods to configure, run, and visualize sweeps of the local learning coefficient over epsilon and beta.

configure_sweep(llc_estimator: Callable, llc_estimator_kwargs: dict, min_epsilon: float | None = 1e-06, max_epsilon: float | None = 0.01, epsilon_samples: float = 8, epsilon_range: List[float] | None = None, min_beta: float | None = None, max_beta: float | None = None, beta_samples: float = 8, beta_range: List[float] | None = None, dataloader: DataLoader | None = None) → None

Configure the sampling parameters for the LLC analysis. :param llc_estimator: Callable function to estimate the local learning coefficient.

Note: The estimator function expected by EpsilonBetaAnalyzer must have the following signature: def estimator(epsilon: float, beta: float, **kwargs) -> dict where kwargs are the arguments to estimate_learning_coeff_with_summary The return value must be a dict with a “llc/trace” key corresponding to a numpy array of shape (num_chains, num_draws) Additional keys can represent other values of interest (e.g. acceptance rates, true LLC.)

Parameters:

  • llc_estimator_kwargs – Keyword arguments for the llc_estimator function.

  • min_epsilon – Minimum value for epsilon range (if epsilon_range not provided).

  • max_epsilon – Maximum value for epsilon range (if epsilon_range not provided).

  • epsilon_samples – Number of samples in epsilon range (if epsilon_range not provided).

  • epsilon_range – Explicit range of epsilon values to use (overrides min/max_epsilon).

  • min_beta – Minimum value for beta range (if beta_range not provided).

  • max_beta – Maximum value for beta range (if beta_range not provided).

  • beta_samples – Number of samples in beta range (if beta_range not provided).

  • beta_range – Explicit range of beta values to use (overrides min/max_beta).

  • dataloader – Optional dataloader for calculating optimal beta.

plot(true_lambda: float | int | str | Container | None = None, num_last_steps_to_average: int = 50, color: str | None = None, slider: str | None = None, slider_plane: str | None = False, **kwargs) → Figure

Plot the results of the LLC sweep.

Parameters:

  • true_lambda – True value of lambda for comparison (optional). Can be a scalar, a list, or a string column name of sweep_df. Will plot a horizontal plane at the true_lambda value.

  • num_last_steps_to_average – Number of last steps to average for final LLC value.

  • color – Column name to use for coloring the scatter points.

  • slider – Column name to use for creating a slider in the plot.

  • slider_plane – If True, adds a plane for each slider value.

  • kwargs – Additional keyword arguments to pass to the plotting function. Example: range_color=[0, 0.15] to set the color range.

Returns:

A plotly Figure object containing the LLC sweep visualization.

save_fig(path: str) → None

Save the current figure to a file. :param path: Path to save the figure to.

sweep(add_to_existing=False) → None

Perform the LLC sweep using the configured parameters.

This method runs the LLC estimator for each combination of epsilon and beta values and stores the results in self.sweep_df.

Parameters:

add_to_existing – If True, adds new sweep results to existing ones. If False, replaces existing results. Useful for sweeping over multiple models or datasets.

Module contents