Cover Image

DrJAX: Scalable and Differentiable MapReduce Primitives in JAX

Fuente: arXiv

Saved in:

Bibliographic Details
Main Authors:	Rush, Keith, Charles, Zachary, Garrett, Zachary, Augenstein, Sean, Mitchell, Nicole
Format:	Preprint
Published:	2024
Subjects:	Distributed, Parallel, and Cluster Computing Machine Learning
Online Access:	Acceder al recurso
Tags:	Add Tag No Tags, Be the first to tag this record!

Similar Items

Federated Automatic Differentiation
by: Rush, Keith, et al.
Published: (2023)

Orbax: Distributed Checkpointing with JAX
by: Gaffney, Colin, et al.
Published: (2026)

Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo
by: Charles, Zachary, et al.
Published: (2025)

Fine-Tuning Large Language Models with User-Level Differential Privacy
by: Charles, Zachary, et al.
Published: (2024)

Scalable Distributed Algorithms for Size-Constrained Submodular Maximization in the MapReduce and Adaptive Complexity Models
by: Chen, Yixin, et al.
Published: (2022)

Scalable Fault-Tolerant MapReduce
by: Hespe, Demian, et al.
Published: (2024)

JORA: JAX Tensor-Parallel LoRA Library for Retrieval Augmented Fine-Tuning
by: Tahir, Anique, et al.
Published: (2024)

JAXMg: A multi-GPU linear solver in JAX
by: Wiersema, Roeland
Published: (2026)

Scalable mRMR feature selection to handle high dimensional datasets: Vertical partitioning based Iterative MapReduce framework
by: Vivek, Yelleti, et al.
Published: (2022)

Wireless MapReduce Arrays for Coded Distributed Computing
by: Peter, Elizabath, et al.
Published: (2024)

Word Frequency Counting Based on Serverless MapReduce
by: Li, Hanzhe, et al.
Published: (2026)

A Novel Approach to Translate Structural Aggregation Queries to MapReduce Code
by: Abdelmoniem, Ahmed M., et al.
Published: (2025)

A New Interference-Alignment Scheme for Wireless MapReduce
by: Bi, Yue, et al.
Published: (2022)

Efficient Candidate-Free R-S Set Similarity Joins with Filter-and-Verification Trees on MapReduce
by: Feng, Yuhong, et al.
Published: (2025)

ADP-VRSGP: Decentralized Learning with Adaptive Differential Privacy via Variance-Reduced Stochastic Gradient Push
by: Wu, Xiaoming, et al.
Published: (2025)

Comparative Evaluation of PyTorch, JAX, SciPy, and Neal for Solving QUBO Problems at Scale
by: Yang, Pei-Kun
Published: (2025)

ScalableHD: Scalable and High-Throughput Hyperdimensional Computing Inference on Multi-Core CPUs
by: Parikh, Dhruv, et al.
Published: (2025)

Efficient AllReduce with Stragglers
by: Devraj, Arjun, et al.
Published: (2025)

Revisiting Reliability in Large-Scale Machine Learning Research Clusters
by: Kokolis, Apostolos, et al.
Published: (2024)

Reducing Energy Bloat in Large Model Training
by: Chung, Jae-Won, et al.
Published: (2023)

Scalable Decentralized Algorithms for Online Personalized Mean Estimation
by: Galante, Franco, et al.
Published: (2024)

Blockchain-based Framework for Scalable and Incentivized Federated Learning
by: Wu, Bijun, et al.
Published: (2025)

Is Flash Attention Stable?
by: Golden, Alicia, et al.
Published: (2024)

Reducing Communication for Split Learning by Randomized Top-k Sparsification
by: Zheng, Fei, et al.
Published: (2023)

Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training
by: Lu, Yishun, et al.
Published: (2026)

Optimizing Federated Learning for Scalable Power-demand Forecasting in Microgrids
by: Banerjee, Roopkatha, et al.
Published: (2025)

JAX-LaB: A High-Performance, Differentiable, Lattice Boltzmann Library for Modeling Multiphase Fluid Dynamics in Geosciences and Engineering
by: Pradhan, Piyush, et al.
Published: (2025)

LLMBridge: Reducing Costs to Access LLMs in a Prompt-Centric Internet
by: Martin, Noah, et al.
Published: (2024)

Parallel Track Transformers: Enabling Fast GPU Inference with Reduced Synchronization
by: Wang, Chong, et al.
Published: (2026)

Federated Learning Framework for Scalable AI in Heterogeneous HPC and Cloud Environments
by: Ghimire, Sangam, et al.
Published: (2025)

Grappa: Gradient-Only Communication for Scalable Graph Neural Network Training
by: Xu, Chongyang, et al.
Published: (2026)

Scalable High-Dimensional Multivariate Linear Regression for Feature-Distributed Data
by: Huang, Shuo-Chieh, et al.
Published: (2023)

A Scalable Digital Twin Framework for Energy Optimization in Data Centers
by: Gonçalves, Raphael Hendrigo de Souza, et al.
Published: (2026)

Reducing Memory Contention and I/O Congestion for Disk-based GNN Training
by: Jiang, Qisheng, et al.
Published: (2024)

An All-Reduce Compatible Top-K Compressor for Communication-Efficient Distributed Learning
by: Chen, Chuyan, et al.
Published: (2025)

Preserving Near-Optimal Gradient Sparsification Cost for Scalable Distributed Deep Learning
by: Yoon, Daegun, et al.
Published: (2024)

High-Dimensional Distributed Sparse Classification with Scalable Communication-Efficient Global Updates
by: Lu, Fred, et al.
Published: (2024)

StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream Generation
by: Zhong, Yinmin, et al.
Published: (2025)

Scalable and Cost-Efficient ML Inference: Parallel Batch Processing with Serverless Functions
by: Barrak, Amine, et al.
Published: (2025)

CycleSL: Server-Client Cyclical Update Driven Scalable Split Learning
by: Wang, Mengdi, et al.
Published: (2025)