cherimoya.cherimoya

The top-level module exposes the model and the EMA helper used during training. The Cheri Block, the fused conv+norm dispatcher, and the inference megakernel internals are in cherimoya.cheri.

Cherimoya

class cherimoya.cherimoya.Cherimoya(*args, **kwargs)

Bases: Module

The Cherimoya sequence-to-function model.

Parameters:

• n_filters (int, optional) – Width of the convolutional backbone (the channel dimension). Default is 128.

• n_layers (int, optional) – Number of stacked Cheri Blocks. Block i uses dilation 2**i. Default is 9.

• signal_groups (list of int, optional) – The number of channels in each signal group. A signal group is one biological modality whose channels share an orientation: a single-channel (unstranded) track is a group of size 1, a stranded (+, -) pair is a group of size 2. The profile head emits one channel per signal channel (total sum(signal_groups) outputs); the count head emits one prediction per group (total len(signal_groups)). Default is [1] — a single unstranded track.

• n_control_tracks (int, optional) – Number of control input tracks (the total channel count summed across all control groups, if any). If 0, the model takes only the one-hot sequence as input. Default is 0.

• expansion (int, optional) – Channel-expansion factor for the MLP inside each Cheri Block. The inner projection maps n_filters -> expansion * n_filters and then back. Default is 2.

• residual_scale (float, optional) – Fixed scalar applied to the MLP output of each Cheri Block before it is added back to the residual stream. Default is 0.15.

• name (str or None, optional) – Display name used when saving model files. Defaults to "cherimoya.{n_filters}.{n_layers}".

• trimming (int or None, optional) – Number of base pairs to trim from each side of the input when producing the output profile. If None, defaults to 46 + sum(2**i for i in range(n_layers)).

• verbose (bool, optional) – Whether the training-progress logger prints to stdout. Default is True.

Constructor

__init__(n_filters=128, n_layers=9, signal_groups=None, n_control_tracks=0, expansion=2, residual_scale=0.15, name=None, trimming=None, verbose=True, compile=True, compile_mode='max-autotune')

save(path)

Save the model to a file.

The checkpoint stores the constructor arguments needed to rebuild the model along with its parameter state dict. This format can be loaded with weights_only=True and is robust to changes in source layout.

Parameters:

path (str) – The destination file path.

classmethod load(path, device='cpu', compile=True, compile_mode='max-autotune')

Load a model previously saved with save().

Parameters:

• path (str) – The checkpoint file path.

• device (str or torch.device, optional) – Device to map the parameters onto. Default is 'cpu'.

• compile (bool, optional) – Whether the loaded model should wrap its forward in torch.compile. Default is True (matches pre-2026-05 behavior). Pass False to get an eager forward — useful for scripts that hit the cudagraph cache-overwrite error or that need to debug / trace the model.

• compile_mode (str, optional) –

  The mode passed through to torch.compile when compile=True. Default is 'max-autotune'. Common alternatives:

  • 'max-autotune-no-cudagraphs' — same kernel autotuning, but disables CUDA graph capture. The safe choice if you hit a cudagraph error but still want autotuned kernels.

  • 'reduce-overhead' — lighter compile, smaller speedup, no autotune sweep.

  Ignored when compile=False.

Returns:

model – The reconstructed model, placed on device.

Return type:

Cherimoya

forward(X, X_ctl=None)

A forward pass of the model.

Dispatches to self._forward_fn (which is either the compiled or eager forward, set in __init__ according to the compile kwarg). Kept as a class-level method so that subclasses overriding forward can still call super().forward(...).

fit(training_data, muon_optimizer, adam_optimizer, lw_optimizer, muon_scheduler, adam_scheduler, lw_scheduler, X_valid, X_ctl_valid, y_valid, max_epochs=50, batch_size=64, dtype='float32', device='cuda', early_stopping=None)

Fit the model to data and validate it periodically.

This method controls the training of a Cherimoya model. It will fit the model to examples generated by the training_data DataLoader object and, if validation data is provided, will validate the model against it at the end of each epoch and return those values.

Two versions of the model will be saved using save(): the best model found during training according to the validation measures, and the final model at the end of training. Additionally, a log will be saved of the training and validation statistics, e.g. time and performance.

Parameters:

• training_data (torch.utils.data.DataLoader) – A generator that produces examples to train on. If n_control_tracks is greater than 0, must product two inputs, otherwise must produce only one input.

• muon_optimizer (torch.optim.Optimizer) – A Muon optimizer to control the training of the 2D non-head/non-tail layers in the model. This is mostly the dense layers and depth-wise convolutions of the Cheri blocks.

• adam_optimizer (torch.optim.Optimizer) – An Adam/W optimizer to control the training of the other parametrers. This should be the head/tail layers, the bias terms, the per-block conv_weight parameter, and any other parameters that are not 2D matrices routed to Muon.

• lw_optimizer (torch.optim.Optimizer) – An optimizer for the Kendall uncertainty weights (lw0, lw1). Typically SGD with momentum.

• muon_scheduler (torch.optim.lr_scheduler) – The scheduler to use for the Muon optimizer. This should likely be a cosine decay with a warmup phase.

• adam_scheduler (torch.optim.lr_scheduler) – The scheduler to use for the Adam/W optimizer. This should likely be the same cosine decay with a warmup phase used for the Muon optimizer.

• lw_scheduler (torch.optim.lr_scheduler) – The scheduler to use for the lw optimizer. Typically a linear warmup followed by a constant rate (no decay).

• X_valid (torch.tensor, shape=(n, 4, length)) – A block of sequences to validate on at the end of each epoch.

• X_ctl_valid (torch.tensor or None, shape=(n, n_control_tracks, length)) – A block of control sequences to use for making the validation set predictions at the end of each epoch. If n_control_tracks is None, pass in None. Default is None.

• y_valid (torch.tensor or None, shape=(n, sum(signal_groups), output_length)) – A block of signals to validate against at the end of each epochs.

• max_epochs (int) – The maximum number of epochs to train for, as measured by the number of times that training_data is exhausted. Default is 50.

• batch_size (int, optional) – The number of examples to include in each batch. Default is 64.

• dtype (str or torch.dtype) – The torch.dtype to use when training. Usually, this will be torch.float32 or torch.bfloat16. Default is torch.float32.

• device (str) – The device to use for training and inference. Typically, this will be ‘cuda’ but can be anything supported by torch. Default is ‘cuda’.

• early_stopping (int or None, optional) – Whether to stop training early. If None, continue training until max_epochs is reached. If an integer, continue training until that number of epochs has been hit without improvement in performance. Default is None.

EMA

class cherimoya.cherimoya.EMA(model, decay=0.999)

Exponential moving average of a model’s parameters.

Maintains a shadow copy of every floating-point parameter that is updated as shadow = decay * shadow + (1 - decay) * parameter after each training step. The shadow weights are typically used at evaluation time, where they tend to produce smoother and more stable predictions than the raw running weights.

Typical usage during training:

1. Create an EMA wrapper after the model is constructed.

2. Call update() after every optimizer step.

3. Call apply_shadow() before evaluation to swap the shadow weights into the model.

4. Call restore() after evaluation to put the training weights back.

Parameters:

• model (torch.nn.Module) – The model whose parameters will be tracked.

• decay (float, optional) – The decay factor of the moving average. Larger values place more weight on the running shadow and less on each new update. Default is 0.999.

Constructor

__init__(model, decay=0.999)

update(model)

Update the shadow weights using the current model parameters.

apply_shadow(model)

Swap the model’s parameters with the shadow weights.

The original weights are kept in an internal backup so they can be restored after evaluation. Calling this method twice in a row without an intervening restore() is an error.

restore(model)

Put the original training weights back into the model.

Used internally by Cherimoya.fit(): a shadow copy of every floating-point parameter (decay 0.999 by default) is updated after every optimizer step, swapped in for validation (apply_shadow/restore), and applied to the saved checkpoints at the end of each improvement and at the very end of training.

The shadow weights are kept on the same device as the model. They are not part of the state_dict and are not saved with Cherimoya.save; they exist only for the duration of the training run.