peal.sparse_dictionaries.batch_topk_network.ActivationStore

class peal.sparse_dictionaries.batch_topk_network.ActivationStore[source]

Bases: object

Minimal next_batch() source over an in-memory activation tensor.

The trainer pulls batches from an object with a next_batch method; this adapter serves consecutive slices of X and wraps around at the end.

Parameters:

  • config (dict or object) – Where batch_size (or model_batch_size) and device are read from; a pydantic config is inspected through its cfg dict.

  • X (torch.Tensor) – Activations of shape (n, act_size).

Methods

next_batch

Return the next (batch_size, act_size) slice moved to device.
__init__(config, X)[source]
next_batch()[source]

Return the next (batch_size, act_size) slice moved to device.