Continuous Thought Machine (CTM). Technical report: https://arxiv.org/abs/2505.05522 Interactive Website: https://pub.sakana.ai/ctm/

Identity Module. Returns the input tensor unchanged. Useful as a placeholder or a no-op layer in nn.Sequential containers or conditional

A custom dataset class that extends the ImageFolder class. Args: root (string): Root directory path. transform (callable, op

LSTM Baseline Args: iterations (int): Number of internal 'thought' steps (T, in paper). d_model (int): Core dimensionality o

LSTM Baseline. Wrapper that lets us use the same backbone as the CTM and LSTM baselines, with a Args: d_model (int): workarou

Squeeze Module. Removes a specified dimension of size 1 from the input tensor. Useful for incorporating tensor dimension squeezing withi

Learnable Fourier Feature Positional Encoding. Implements Algorithm 1 from "Learnable Fourier Features for Multi-Dimensional Spatial Pos

SuperLinear Layer: Implements Neuron-Level Models (NLMs) for the CTM. This layer is the core component enabling Neuron-Level Models (NLMs),

Custom Rotational Positional Embedding. Generates 2D positional embeddings based on rotating a fixed start vector. The rotation angle fo

r""" A distributed sampler that continuously yields random indices with replacement, avoiding frequent iterator recreation overhead for DataLo

Simple backbone for MNIST feature extraction.

Combines multiple LearnableFourierPositionalEncoding modules with different initialization scales (gamma) via a learnable weighted sum.

A QAMNIST dataset that includes plus and minus operations on MNIST digits.

Simple, wide, shallow convolutional backbone for image feature extraction. Alternative to ResNet, uses grouped convolutions and GLU activati

UNET-style architecture for the Synapse Model (f_theta1 in the paper). This module implements the connections between neurons in the CTM's l

Slight adaption of the CTM to work with the sort task.

LSTM Baseline Args: iterations (int): Number of internal 'thought' steps (T, in paper). d_model (int): Core dimensionality o

LSTM Baseline Args: iterations (int): Number of internal 'thought' steps (T, in paper). d_model (int): Core dimensionality

Simple wrapper that automatically resets the environment on done. Modeled after EpisodicLifeEnv but simplified since we don't need to han

Helper class for managing "thought steps" in the ctm_qamnist pipeline. Args: iterations_per_digit (int): Number of iterations for ea

Wrapper to use standard pre-trained ResNet models from torchvision. Loads a specified ResNet architecture pre-trained on ImageNet, removes t