cherimoya 0.0.1__tar.gz

cherimoya-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jacob Schreiber
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

cherimoya-0.0.1/PKG-INFO

@@ -0,0 +1,28 @@
+Metadata-Version: 2.4
+Name: cherimoya
+Version: 0.0.1
+Summary: A light-weight deep learning architecture that uses modern optimization tricks to achieve strong predictive performance.
+Home-page: https://github.com/jmschrei/cherimoya
+Author: Jacob Schreiber
+Author-email: jmschreiber91@gmail.com
+License: LICENSE.txt
+License-File: LICENSE
+Requires-Dist: numpy>=1.14.2
+Requires-Dist: scipy>=1.0.0
+Requires-Dist: pandas>=1.3.3
+Requires-Dist: torch>=1.9.0
+Requires-Dist: h5py>=3.7.0
+Requires-Dist: tqdm>=4.64.1
+Requires-Dist: seaborn>=0.11.2
+Requires-Dist: modisco>=2.0.0
+Requires-Dist: tangermeme>=0.2.3
+Requires-Dist: macs3
+Requires-Dist: bam2bw
+Requires-Dist: bpnet-lite
+Dynamic: author
+Dynamic: author-email
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: summary

cherimoya-0.0.1/README.md

@@ -0,0 +1,40 @@
+<img src="https://github.com/jmschrei/cherimoya/blob/main/imgs/cherimoya.png">
+
+> [!IMPORTANT]
+> Cherimoya is still under active development and may change in ways that are not back compatible. Please make note of the version you are using in case you need to return to it in the future.
+
+Cherimoya is a lightweight genomic sequence-to-function (S2F) model for predicting genomic modalities such as transcription factor binding, chromatin accessibility, and transcription initiation. It builds on concepts that were first introduced by BPNet and ChromBPNet while introducing architectural, algorithmic, and systems-level improvements that improve training stability, efficiency, and predictive performance. Despite needing significantly fewer parameters than other architectures, Cherimoya achieves strong predictive performance across a range of tasks and runs ~5-15x faster when measured on an H200 GPU. 
+
+<img src="https://github.com/jmschrei/cherimoya/blob/main/imgs/cheri-model.png">
+
+The secret to Cherimoya's success is a new Cheri Block, which adapts the ConvNeXT block to the domain of noisy high-throughput genomics experiments. This block is comprised of a dilated depth-wise convolution, a layer norm, a projection into a higher-dimensional space, a GeLU non-linearity, a projection back into the original dimensionality, and then a channel-wise scaling for robustness. Conceptually, this means that the blocks first aggregate information spatially but independently for each feature/channel (the depth-wise convolution) and then aggregate information across features but independently for each position (the two projections). The dilated depth-wise convolution and the layer norm have been fused into an efficient custom GPU kernel that is ~2-3x faster than the native PyTorch implementation.
+
+<img src="https://github.com/jmschrei/cherimoya/blob/main/imgs/cheri-block.png">
+
+---
+
+## Key Features
+
+### Lightweight Architecture
+Cherimoya employs a compact convolutional backbone that substantially reduces parameter count without sacrificing predictive accuracy. This design enables efficient training, large-scale hyperparameter exploration, interactive usage via browsers, and usage of dozens or hundreds of such models simultaneously in complex design settings.
+
+### Automatic Loss Weight Balancing
+Profile and count losses are combined using learned weighting parameters rather than fixed hyperparameters. This approach replaces the heuristic developed for BPNet and ChromBPNet models and enables the models to scale to larger contexts and across modalities automatically, while also improving gradient stability across datasets with varying signal-to-noise characteristics.
+
+### Muon Optimizer
+Cherimoya uses the Muon optimizer when training the projection layers, and the AdamW optimizer for all other layers. This has significantly accelerated training by reducing the number of epochs needed while modestly improving performance.
+
+### End-to-End Pipeline
+Cherimoya provides an integrated pipeline covering:
+
+- BAM/SAM/fragment file conversion using bam2bw
+- Peak calling using MACS3
+- model training
+- evaluation
+- downstream analysis and motif discovery using TF-MoDISco
+
+<img src="https://github.com/jmschrei/cherimoya/blob/main/imgs/pipeline.png">
+  
+
+This design supports reproducible end-to-end experiments and reduces the overhead associated with managing separate tooling for each stage.
+

cherimoya-0.0.1/cherimoya/__init__.py

@@ -0,0 +1,6 @@
+# cherimoya
+# Author: Jacob Schreiber
+
+from .cherimoya import Cherimoya
+
+__version__ = '0.0.1'

cherimoya-0.0.1/cherimoya/cherimoya.py

@@ -0,0 +1,604 @@
+# cherimoya.py
+# Author: Jacob Schreiber <jmschreiber91@gmail.com>
+
+"""
+An implementing of the Cherimoya deep learning model, which is a compact
+architecture for predicting genomic modalities from sequence alone.
+"""
+
+import h5py
+import time 
+import numpy
+import torch
+
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+import itertools
+
+from .losses import MNLLLoss
+from .losses import log1pMSELoss
+from .losses import _mixture_loss
+
+from .performance import pearson_corr
+from .performance import calculate_performance_measures
+
+from tqdm import tqdm
+
+from tangermeme.predict import predict
+from bpnetlite.logging import Logger
+
+torch.backends.cudnn.benchmark = True	
+torch.set_float32_matmul_precision('high')
+
+
+def autotune_configs():
+    num_warps = [4, 8, 16]
+    num_stages = [2, 3, 4, 5]
+    BLOCK_Ls = [32, 64, 128, 256]
+    
+    configs = []
+    for num_warp, num_stage, L in itertools.product(num_warps, num_stages, BLOCK_Ls):
+        configs.append(triton.Config({
+            'num_warps': num_warp,
+            'num_stages': num_stage,
+            'BLOCK_L': L
+        }))
+    return configs
+
+
+@triton.autotune(
+    configs = autotune_configs(),
+    key=['C', 'L'],
+)
+@triton.jit
+def fwd_conv_kernel(
+    X_ptr, W_ptr, Y_ptr, Mean_ptr, Rstd_ptr,
+    stride_xn, dilation, eps,
+	L: tl.constexpr,
+	C: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+    BLOCK_L: tl.constexpr
+):
+	
+	pid_n = tl.program_id(0)
+	offs_c = tl.arange(0, BLOCK_C)[None, :]
+	mask_c = offs_c < C
+	
+	w_idx = W_ptr + offs_c
+	w0 = tl.load(w_idx,       mask=mask_c, other=0.0).to(tl.float32)
+	w1 = tl.load(w_idx + C,   mask=mask_c, other=0.0).to(tl.float32)
+	w2 = tl.load(w_idx + C*2, mask=mask_c, other=0.0).to(tl.float32)
+	
+	running_sum = 0.0
+	running_sq_sum = 0.0
+	for l_start in tl.range(0, L, BLOCK_L):
+		offs = l_start + tl.arange(0, BLOCK_L)[:, None]
+		offs_l = offs - dilation
+		offs_r = offs + dilation
+		
+		mask = (offs < L) & mask_c
+		mask_l = (offs_l >= 0) & mask
+		mask_r = (offs_r < L) & mask
+		
+		x_idx = X_ptr + pid_n * stride_xn + offs_c
+		x1 = tl.load(x_idx + offs*C,   mask=mask,   other=0.0).to(tl.float32)
+		x0 = tl.load(x_idx + offs_l*C, mask=mask_l, other=0.0).to(tl.float32)
+		x2 = tl.load(x_idx + offs_r*C, mask=mask_r, other=0.0).to(tl.float32)
+		
+		conv = x0*w0 + x1*w1 + x2*w2
+		conv2 = conv * conv
+	
+		running_sum += tl.sum(conv)
+		running_sq_sum += tl.sum(conv2)
+	
+	count = L * C
+	mean = running_sum / count
+	var = (running_sq_sum / count) - (mean * mean)
+	rstd = 1.0 / tl.sqrt(var + eps)
+	
+	tl.store(Mean_ptr + pid_n, mean)
+	tl.store(Rstd_ptr + pid_n, rstd)
+	
+	for l_start in tl.range(0, L, BLOCK_L):
+		offs = l_start + tl.arange(0, BLOCK_L)[:, None]
+		offs_l = offs - dilation
+		offs_r = offs + dilation
+		
+		mask = (offs < L) & mask_c
+		mask_l = (offs_l >= 0) & mask
+		mask_r = (offs_r < L) & mask
+		
+		x_idx = X_ptr + pid_n * stride_xn + offs_c
+		x1 = tl.load(x_idx + offs*C,   mask=mask, other=0.0).to(tl.float32)
+		x0 = tl.load(x_idx + offs_l*C, mask=mask_l, other=0.0).to(tl.float32)
+		x2 = tl.load(x_idx + offs_r*C, mask=mask_r, other=0.0).to(tl.float32)
+	
+		conv = x0*w0 + x1*w1 + x2*w2
+		x_hat = (conv - mean) * rstd
+	
+		y_idx = Y_ptr + pid_n * stride_xn + offs * C + offs_c
+		tl.store(y_idx, x_hat, mask=mask)
+
+
+@triton.autotune(
+    configs = autotune_configs(),
+    key=['C', 'L']
+)
+@triton.jit
+def bwd_conv_kernel(
+	dY_ptr, X_ptr, W_ptr, Mean_ptr, Rstd_ptr,
+    dX_ptr, dW_ptr, stride_xn, dilation,
+	L: tl.constexpr, 
+	C: tl.constexpr, 
+	BLOCK_C: tl.constexpr, 
+	BLOCK_L: tl.constexpr
+):
+	pid_n = tl.program_id(0)
+	offs_c = tl.arange(0, BLOCK_C)
+	mask_c = offs_c < C
+	
+	mean = tl.load(Mean_ptr + pid_n)
+	rstd = tl.load(Rstd_ptr + pid_n)
+	
+	w_idx = W_ptr + offs_c
+	w0 = tl.load(w_idx,       mask=mask_c)[None, :].to(tl.float32)
+	w1 = tl.load(w_idx + C,   mask=mask_c)[None, :].to(tl.float32)
+	w2 = tl.load(w_idx + C*2, mask=mask_c)[None, :].to(tl.float32)
+	
+	sum_dy = 0.0
+	sum_dy_xhat = 0.0
+	
+	mask_c = mask_c[None, :]
+	offs_c = offs_c[None, :]
+	
+	for l_start in tl.range(0, L, BLOCK_L):
+		offs = l_start + tl.arange(0, BLOCK_L)[:, None]
+		offs_l = offs - dilation
+		offs_r = offs + dilation
+		
+		mask = (offs < L) & mask_c
+		mask_l = (offs_l >= 0) & mask
+		mask_r = (offs_r < L) & mask
+		
+		x_idx = X_ptr + pid_n * stride_xn + offs_c
+		x1 = tl.load(x_idx + offs*C,   mask=mask, other=0.0).to(tl.float32)
+		x0 = tl.load(x_idx + offs_l*C, mask=mask_l, other=0.0).to(tl.float32)
+		x2 = tl.load(x_idx + offs_r*C, mask=mask_r, other=0.0).to(tl.float32)
+		
+		conv = x0*w0 + x1*w1 + x2*w2
+		x_hat = (conv - mean) * rstd
+		
+		y_idx = dY_ptr + pid_n * stride_xn + offs * C + offs_c
+		dy = tl.load(y_idx, mask=mask, other=0.0).to(tl.float32)
+		
+		sum_dy += tl.sum(dy)
+		sum_dy_xhat += tl.sum(dy * x_hat)
+	
+	dw0 = tl.zeros((1, BLOCK_C), dtype=tl.float32)
+	dw1 = tl.zeros((1, BLOCK_C), dtype=tl.float32)
+	dw2 = tl.zeros((1, BLOCK_C), dtype=tl.float32)
+	
+	for l_start in tl.range(0, L, BLOCK_L):
+		offs = l_start + tl.arange(0, BLOCK_L)[:, None]
+		offs_l = offs - dilation
+		offs_r = offs + dilation
+		
+		mask = (offs < L) & mask_c
+		mask_l = (offs_l >= 0) & mask
+		mask_r = (offs_r < L) & mask
+		
+		x_idx = X_ptr + pid_n * stride_xn + offs_c
+		x1 = tl.load(x_idx + offs*C,   mask=mask, other=0.0).to(tl.float32)
+		x0 = tl.load(x_idx + offs_l*C, mask=mask_l, other=0.0).to(tl.float32)
+		x2 = tl.load(x_idx + offs_r*C, mask=mask_r, other=0.0).to(tl.float32)
+		
+		conv = x0*w0 + x1*w1 + x2*w2
+		x_hat = (conv - mean) * rstd
+	
+		###
+		
+		dy_idx = dY_ptr + pid_n * stride_xn + offs * C + offs_c
+		dy = tl.load(dy_idx, mask=mask, other=0.0).to(tl.float32)
+	
+		count = L * C
+		d_conv = (rstd / count) * (count * dy - sum_dy - x_hat * sum_dy_xhat)
+		
+		dw0 += tl.sum(d_conv * x0, axis=0)[None, :]
+		dw1 += tl.sum(d_conv * x1, axis=0)[None, :]
+		dw2 += tl.sum(d_conv * x2, axis=0)[None, :]
+	
+		dx_idx0 = dX_ptr + pid_n * stride_xn + offs * C + offs_c
+		
+		dx1 = tl.load(dx_idx0, mask=mask, other=0.0)
+		dx0 = tl.load(dx_idx0 - dilation*C, mask=mask_l, other=0.0)
+		dx2 = tl.load(dx_idx0 + dilation*C, mask=mask_r, other=0.0)
+	
+		dx1 += d_conv * w1
+		dx0 += d_conv * w0
+		dx2 += d_conv * w2
+	
+		tl.store(dx_idx0,              dx1, mask=mask)
+		tl.store(dx_idx0 - dilation*C, dx0, mask=mask_l)
+		tl.store(dx_idx0 + dilation*C, dx2, mask=mask_r)
+	
+	
+	dw_idx = dW_ptr + pid_n * (C * 3) + offs_c
+	tl.store(dw_idx,       dw0, mask=mask_c)
+	tl.store(dw_idx + C,   dw1, mask=mask_c)
+	tl.store(dw_idx + C*2, dw2, mask=mask_c)
+
+
+class FusedDilatedConvNormFunc(torch.autograd.Function):
+	@staticmethod
+	def forward(ctx, x, w, dilation):
+		N, L, C = x.shape
+		BLOCK_C = triton.next_power_of_2(C)
+		eps = 1e-3
+		
+		mean = torch.empty((N,), dtype=torch.float32, device=x.device)
+		rstd = torch.empty((N,), dtype=torch.float32, device=x.device)
+		y = torch.empty_like(x)
+		
+		fwd_conv_kernel[(N,)](
+			x, w, y, mean, rstd,
+			x.stride(0), dilation, eps,
+			L, C, BLOCK_C=BLOCK_C
+		)
+	
+		ctx.save_for_backward(x, w, mean, rstd)
+		ctx.dilation = dilation
+		return y
+	
+	@staticmethod
+	def backward(ctx, dy):
+		x, w, mean, rstd = ctx.saved_tensors
+		N, L, C = x.shape
+		BLOCK_C = triton.next_power_of_2(C)
+		
+		dy = dy.contiguous()
+		dx = torch.zeros_like(x, dtype=x.dtype)
+		dw = torch.empty((N, 3, C), device=x.device, dtype=torch.float32)
+		
+		bwd_conv_kernel[(N,)](
+			dy, x, w, mean, rstd, dx, dw, 
+			x.stride(0), ctx.dilation,
+			L, C, BLOCK_C, 
+		)
+	
+		dw = dw.sum(dim=0)		
+		return dx.to(x.dtype), dw.to(x.dtype), None
+
+
+###
+
+		
+class CheriBlock(torch.nn.Module):
+	def __init__(self, n_filters, dilation, eps=0.01):
+		super().__init__()
+		self.n_filters = n_filters
+		self.dilation = dilation
+
+		self.conv_weight = torch.nn.Parameter(torch.randn(3, n_filters))	
+		self.linear1 = torch.nn.Linear(n_filters, 2*n_filters, bias=False)
+		self.linear2 = torch.nn.Linear(2*n_filters, n_filters, bias=False)
+		self.gamma = torch.nn.Parameter(torch.ones(1, n_filters) * eps)
+		self.activation = torch.nn.GELU(approximate='tanh')
+
+		torch.nn.init.trunc_normal_(self.conv_weight, std=0.02)
+		torch.nn.init.trunc_normal_(self.linear1.weight, std=0.02)
+		torch.nn.init.trunc_normal_(self.linear2.weight, std=0.02)
+	
+	def forward(self, X):
+		X_conv = FusedDilatedConvNormFunc.apply(X, self.conv_weight, self.dilation)
+		X_mlp = self.linear2(self.activation(self.linear1(X_conv)))
+		return X + X_mlp * self.gamma
+		
+
+class CheriBlock2(torch.nn.Module):
+	def __init__(self, n_filters, dilation, eps=0.01):
+		super().__init__()
+		self.n_filters = n_filters
+		self.dilation = dilation
+
+		self.conv = torch.nn.Conv1d(n_filters, n_filters, groups=n_filters, dilation=dilation, padding=dilation, kernel_size=3)
+		#self.norm = torch.nn.LayerNorm((n_filters, 2114), elementwise_affine=False, bias=False, eps=1e-3)		
+		self.linear1 = torch.nn.Conv1d(n_filters, 3*n_filters, kernel_size=1, bias=False)
+		self.linear2 = torch.nn.Conv1d(3*n_filters, n_filters, kernel_size=1, bias=False)
+		self.gamma = torch.nn.Parameter(torch.ones(n_filters, 1) * eps) 
+		self.activation = torch.nn.GELU(approximate='tanh')
+		
+		torch.nn.init.trunc_normal_(self.conv.weight, std=0.02)
+		torch.nn.init.trunc_normal_(self.linear1.weight, std=0.02)
+		torch.nn.init.trunc_normal_(self.linear2.weight, std=0.02)
+	
+	def forward(self, X):
+		X_conv = self.conv(X)
+		X_conv = self.norm(X_conv)
+		X_mlp = self.linear2(self.activation(self.linear1(X_conv)))
+		X = torch.add(X, X_mlp * self.gamma)
+		return X
+
+
+
+class Cherimoya(torch.nn.Module):
+	def __init__(self, n_filters=64, n_layers=9, n_outputs=1, 
+		n_control_tracks=0, name=None, trimming=None, 
+		single_count_output=True, verbose=True):
+		super(Cherimoya, self).__init__()
+		self.n_filters = n_filters
+		self.n_layers = n_layers
+		self.n_outputs = n_outputs
+		self.n_control_tracks = n_control_tracks
+
+		self.name = name or "cherimoya.{}.{}".format(n_filters, n_layers)
+		self.trimming = trimming or 46 + sum(2**i for i in range(n_layers))
+
+		self.iconv = torch.nn.Conv1d(4, n_filters, kernel_size=19, padding=9)
+		self.igelu = torch.nn.GELU(approximate='tanh')
+
+		self.blocks = torch.nn.ModuleList([
+			CheriBlock(n_filters, 2**i) for i in range(self.n_layers)
+		])
+		
+		self.fconv = torch.nn.Conv1d(n_filters+n_control_tracks, n_outputs, 
+			kernel_size=75, padding=37)
+
+		self.lw0 = torch.nn.Parameter(torch.ones(1)) 
+		self.lw1 = torch.nn.Parameter(torch.ones(1))
+
+		n_count_control = 1 if n_control_tracks > 0 else 0
+		n_count_outputs = 1 if single_count_output else n_outputs
+		self.linear = torch.nn.Linear(n_filters+n_count_control, n_count_outputs)
+		
+		torch.nn.init.trunc_normal_(self.iconv.weight, std=0.02)
+		torch.nn.init.trunc_normal_(self.fconv.weight, std=0.02)
+		torch.nn.init.trunc_normal_(self.linear.weight, std=0.02)
+
+		torch.nn.init.zeros_(self.iconv.bias)
+		torch.nn.init.zeros_(self.fconv.bias)
+		torch.nn.init.zeros_(self.linear.bias)
+
+		self.logger = Logger(["Epoch", "Iteration", "Training Time",
+			"Validation Time", "Training MNLL", "Training Count MSE", 
+			"Validation MNLL", "Validation Profile Pearson", 
+			"Validation Count Pearson", "Validation Count MSE", "Saved?"], 
+			verbose=verbose)
+
+
+	@torch.compile(mode='max-autotune')
+	def forward(self, X, X_ctl=None):
+		"""A forward pass of the model.
+
+		This method takes in a nucleotide sequence X, a corresponding
+		per-position value from a control track, and a per-locus value
+		from the control track and makes predictions for the profile 
+		and for the counts. This per-locus value is usually the
+		log(sum(X_ctl_profile)+1) when the control is an experimental
+		read track but can also be the output from another model.
+
+		Parameters
+		----------
+		X: torch.tensor, shape=(batch_size, 4, length)
+			The one-hot encoded batch of sequences.
+
+		X_ctl: torch.tensor or None, shape=(batch_size, n_strands, length)
+			A value representing the signal of the control at each position in 
+			the sequence. If no controls, pass in None. Default is None.
+
+		Returns
+		-------
+		y_profile: torch.tensor, shape=(batch_size, n_strands, out_length)
+			The output predictions for each strand trimmed to the output
+			length.
+		"""
+
+		start, end = self.trimming, X.shape[2] - self.trimming
+		
+		X = self.igelu(self.iconv(X))
+		X = X.transpose(1, 2).contiguous()
+		for i in range(self.n_layers):
+			X = self.blocks[i](X)
+
+		X = X.transpose(1, 2).contiguous()
+		if X_ctl is None:
+			X_w_ctl = X
+		else:
+			X_w_ctl = torch.cat([X, X_ctl], dim=1)
+
+		y_profile = self.fconv(X_w_ctl)[:, :, start:end]
+
+		# counts prediction
+		X = torch.mean(X[:, :, start-37:end+37].float(), dim=2)
+		if X_ctl is not None:
+			X_ctl = torch.sum(X_ctl[:, :, start-37:end+37].float(), dim=(1, 2))
+			X_ctl = X_ctl.unsqueeze(-1)
+			X = torch.cat([X, torch.log(X_ctl+1)], dim=-1)
+
+		y_counts = self.linear(X)
+		return y_profile, y_counts
+
+		
+	def fit(self, training_data, muon_optimizer, adam_optimizer, muon_scheduler,
+		adam_scheduler, X_valid, X_ctl_valid, y_valid, max_epochs=100, 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 BPNet 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: 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, and any other parameters
+			that are not 2D matrices.
+
+		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.
+
+		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, n_outputs, 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 100.
+
+		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.
+		"""
+		
+		if X_valid is not None:
+			y_valid_counts = y_valid.sum(dim=2)
+
+		if X_ctl_valid is not None:
+			X_ctl_valid = (X_ctl_valid,)
+			
+		dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype
+
+		iteration = 0
+		early_stop_count = 0
+		best_loss = float("inf")
+		self.logger.start()
+		
+		###
+
+		for epoch in range(max_epochs):
+			tic = time.time()
+			
+			for data in training_data:
+				X, y, labels = data[0], data[-2], data[-1]
+				X_ctl = data[1].to(device) if len(data) == 4 else None
+
+				if X.shape[0] != batch_size:
+					continue
+
+				X = X.to(device).float()
+				y = y.to(device)
+				
+				# Clear the optimizer and set the model to training mode
+				muon_optimizer.zero_grad()
+				adam_optimizer.zero_grad()
+				self.train()
+
+				# Make one training step
+				with torch.autocast(device_type=device, dtype=dtype):
+					y_hat_logits, y_hat_logcounts = self(X, X_ctl)
+
+				profile_loss, count_loss = _mixture_loss(y, y_hat_logits.float(), 
+					y_hat_logcounts.float(), )
+
+
+				w0 = (1.0 / (2.0 * self.lw0 ** 2))
+				w1 = (1.0 / (2.0 * self.lw1 ** 2))
+				loss = w0*profile_loss + w1*count_loss
+				
+				if self.lw0.requires_grad == True:
+					loss += torch.sum(torch.log(self.lw0) ** 2 + torch.log(self.lw1) ** 2)
+				
+				loss.backward()
+				
+				muon_optimizer.step()
+				adam_optimizer.step()
+
+				muon_scheduler.step()
+				adam_scheduler.step()
+
+				iteration += 1
+
+			train_time = time.time() - tic 
+			
+			if self.lw0.requires_grad == True and torch.abs(self.lw0.grad).sum() < 1:
+				self.lw0.requires_grad = False
+				self.lw1.requires_grad = False
+				
+			# Validate the model at the end of the epoch
+			with torch.no_grad():
+				self.eval()
+				tic = time.time()
+
+				y_hat_logits, y_hat_logcounts = predict(self, X_valid, args=X_ctl_valid, 
+					batch_size=batch_size, dtype=dtype, device=device)
+				
+				valid_profile_loss, valid_count_loss = _mixture_loss(y_valid,
+					y_hat_logits, y_hat_logcounts)
+
+				valid_loss = w0*valid_profile_loss + w1*valid_count_loss
+				
+				measures = calculate_performance_measures(y_hat_logits,
+					y_valid, y_hat_logcounts, measures=['profile_pearson', 'count_pearson'])
+
+				valid_profile_corr = numpy.nan_to_num(measures['profile_pearson'])
+				valid_count_corr = numpy.nan_to_num(measures['count_pearson'])
+				valid_time = time.time() - tic
+				
+				self.logger.add([epoch, 
+					iteration, 
+					train_time, 
+					valid_time,
+					profile_loss.item(), 
+					count_loss.item(), 
+					valid_profile_loss.item(), 
+					valid_profile_corr.mean(),
+					valid_count_corr.mean(), 
+					valid_count_loss.item(),
+					(valid_loss < best_loss).item()])
+
+				self.logger.save("{}.log".format(self.name))
+				
+				if valid_loss < best_loss:
+					torch.save(self, "{}.torch".format(self.name))
+					best_loss = valid_loss
+					early_stop_count = -1
+
+			early_stop_count += 1
+			if early_stopping is not None and early_stop_count >= early_stopping:
+				break
+
+		torch.save(self, "{}.final.torch".format(self.name))