Representax: Scalable Representation Learning in JAX

Learned representations shape how models organize information and generalize beyond their training examples. As learning systems evolve to incorporate information from increasingly diverse modalities, a flexible and efficient foundation for representation-learning research becomes essential to enable progress. We introduce Representax, an open-source JAX framework for developing and training representation-learning models across modalities. Representax unifies extensible training and evaluation tasks with established methods spanning dense and late-interaction retrieval, multimodal contrastive learning, reward modeling, and joint-embedding predictive learning. Its shared training infrastructure supports model adaptation, multimodal data processing, memory-efficient optimization, and distributed execution on GPUs and TPUs. To demonstrate its capabilities, we compare training performance against established frameworks, train and evaluate retrieval and multimodal models, and measure multi-GPU strong scaling. These experiments demonstrate competitive throughput, improvements on held-out retrieval tasks, and efficient scaling. Representax provides a foundation for investigating how representations are learned and adapted, and for taking new ideas from exploration to scale.