API Reference#

Look through specific module or search through the entire list of APIs.

All APIs

Object	Description
`Dataset`	PyTorch Dataset for loading time-slices of neural data recordings from HDF5 files. `torch_brain.dataset`
`DatasetIndex`	Index for accessing a specific time interval of a recording within a `Dataset`. `torch_brain.dataset`
`NestedDataset`	Dataset that composes multiple `Dataset` instances under a single interface. `torch_brain.dataset`
`SpikingDatasetMixin`	Mixin class for `torch_brain.dataset.Dataset` subclasses containing spiking data. `torch_brain.dataset`
`CalciumImagingDatasetMixin`	Mixin class for `torch_brain.dataset.Dataset` subclasses containing calcium imaging data. `torch_brain.dataset`
`MultiChannelDatasetMixin`	Mixin class for `torch_brain.dataset.Dataset` subclasses containing `torch_brain.dataset`
`chain`	Wrap an object to specify that it (or any of its members) should be stacked `torch_brain.data`
`collate`	Extension of PyTorch’s `default_collate` function to enable more advanced `torch_brain.data`
`pad`	Wrap an object to specify that it (or any of its members) should be padded to `torch_brain.data`
`pad8`	Wrap an object to specify that it (or any of its members) should be padded to `torch_brain.data`
`pad2d`	type obj: `torch_brain.data`
`pad2d8`	type obj: `torch_brain.data`
`track_batch`	Wrap an array or tensor to track the batch_index. `torch_brain.data`
`track_mask`	Wrap an array or tensor to specify that its padding mask should be tracked. `torch_brain.data`
`track_mask8`	Wrap an array or tensor to specify that its padding mask should be tracked. This `torch_brain.data`
`track_mask2d`	Wrap an array or tensor to specify that its padding mask should be tracked. This `torch_brain.data`
`track_mask2d8`	Wrap an array or tensor to specify that its padding mask should be tracked. This `torch_brain.data`
`Dataset`	Deprecated. Please use `torch_brain.dataset.Dataset`. `torch_brain.data`
`RandomFixedWindowSampler`	Samples fixed-length windows randomly, given intervals defined in the `torch_brain.data.sampler`
`SequentialFixedWindowSampler`	Samples fixed-length windows sequentially, always in the same order. The `torch_brain.data.sampler`
`TrialSampler`	Randomly samples a single trial interval from the given intervals. `torch_brain.data.sampler`
`DistributedEvaluationSamplerWrapper`	Wraps a sampler to be used in a distributed evaluation setting. Unlike the standard `torch_brain.data.sampler`
`DistributedStitchingFixedWindowSampler`	A sampler designed specifically for evaluation that enables sliding window `torch_brain.data.sampler`
`Compose`	Compose several transforms together. All transforms will be called sequentially, `torch_brain.transforms`
`RandomChoice`	Apply a single transformation randomly picked from a list. `torch_brain.transforms`
`ConditionalChoice`	Conditionally apply a single transformation based on whether a condition is met. `torch_brain.transforms`
`UnitDropout`	Augmentation that randomly drops units from the sample. By default, the number `torch_brain.transforms`
`TriangleDistribution`	Triangular distribution with a peak at mode_units, going from min_units to `torch_brain.transforms`
`RandomTimeScaling`	<no summary> `torch_brain.transforms`
`RandomOutputSampler`	<no summary> `torch_brain.transforms`
`RandomCrop`	<no summary> `torch_brain.transforms`
`BinSpikes`	Bin spike events into fixed-width time bins. `torch_brain.transforms`
`UnitFilter`	Drop units based on the mask_fn given in the constructor. `torch_brain.transforms`
`UnitFilterById`	Keep/drop units based on the keyword/regex given in the constructor. `torch_brain.transforms`
`Embedding`	A simple extension of `torch.nn.Embedding` to allow more control over `torch_brain.nn`
`InfiniteVocabEmbedding`	Embedding layer with a vocabulary that can be extended. Vocabulary is saved along `torch_brain.nn`
`RotaryTimeEmbedding`	Rotary time/positional embedding layer. This module is designed to be used with `torch_brain.nn`
`SinusoidalTimeEmbedding`	Sinusoidal time/position embedding layer. `torch_brain.nn`
`FeedForward`	A feed-forward network with GEGLU activation. `torch_brain.nn`
`RotaryCrossAttention`	Cross-attention layer with rotary positional embeddings. `torch_brain.nn`
`RotarySelfAttention`	Self-attention layer with rotary positional embeddings. `torch_brain.nn`
`MultitaskReadout`	A module that performs multi-task linear readouts from output embeddings. `torch_brain.nn`
`prepare_for_multitask_readout`	<no summary> `torch_brain.nn`
`Loss`	Base class for losses. All losses should support an optional weights argument. `torch_brain.nn`
`MSELoss`	Base class for losses. All losses should support an optional weights argument. `torch_brain.nn`
`CrossEntropyLoss`	Base class for losses. All losses should support an optional weights argument. `torch_brain.nn`
`MallowDistanceLoss`	Base class for losses. All losses should support an optional weights argument. `torch_brain.nn`
`POYO`	POYO model from Azabou et al. 2023, A Unified, Scalable Framework for Neural Population Decoding. `torch_brain.models`
`POYOPlus`	POYO+ model from Azabou et al. 2025, Multi-session, multi-task neural decoding from distinct cell-types and brain regions. `torch_brain.models`
`CalciumPOYOPlus`	Calcium POYO+ model from Azabou et al. 2025, Multi-session, multi-task neural decoding from distinct cell-types and brain regions. `torch_brain.models`
`poyo_mp`	<no summary> `torch_brain.models`
`DataType`	Enum defining the possible data types. `torch_brain.registry`
`ModalitySpec`	Specification for a modality. `torch_brain.registry`
`register_modality`	Register a new modality specification in the global registry. `torch_brain.registry`
`get_modality_by_id`	Get a modality specification by its ID. `torch_brain.registry`
`MODALITY_REGISTRY`	dict() -> new empty dictionary `torch_brain.registry`
`stitch`	Pools values that share the same timestamp using mean or mode operations. `torch_brain.utils`
`seed_everything`	Sets random seed for reproducibility. `torch_brain.utils`
`create_linspace_latent_tokens`	Creates a sequence of latent tokens. Each token is defined by the `torch_brain.utils`
`create_start_end_unit_tokens`	Creates for each unit a start and end token. Each token is defined by the `torch_brain.utils`
`resolve_weights_based_on_interval_membership`	Determine weights for timestamps based on which intervals they fall within. `torch_brain.utils`
`isin_interval`	Check if timestamps are in any of the intervals in the Interval object. `torch_brain.utils`
`prepare_for_readout`	<no summary> `torch_brain.utils`
`np_string_prefix`	Adds a string prefix to each element of a numpy string array. `torch_brain.utils`
`bin_spikes`	Bins spikes into time bins of size bin_size. If the total time spanned by `torch_brain.utils`