{"categories":[{"label":"Distributed Computing","url":"https://skillfed.io/packages/category/system-distributed-computing/2"}],"enrichment":{"capability":"DrJAX embeds MapReduce programming primitives into JAX, enabling differentiable distributed and parallel computations with optimized execution across large-scale datacenter settings.","skillfed_tags":["distributed-computing","jax-ecosystem","federated-learning"],"use_cases":["Implement distributed training loops for large models using MapReduce patterns with automatic differentiation.","Build federated learning workflows where computations are distributed across multiple devices or nodes.","Optimize parallel data processing pipelines in datacenters by leveraging JAX's sharding for efficient communication.","Prototype distributed algorithms that require differentiable reductions and broadcasts across compute clusters.","Combine MapReduce-style data aggregation with JAX's automatic differentiation for gradient-based distributed optimization."],"what_it_does":"DrJAX is a JAX library that brings MapReduce-style distributed computing into the JAX ecosystem. It provides a simple authoring surface for parallel computations while leveraging JAX's sharding mechanisms for optimized execution in large-scale settings. The package is designed around three core goals: offering an accessible way to write MapReduce computations in JAX, enabling efficient execution via JAX's built-in parallelization, and maintaining full differentiability throughout\u2014including through communication primitives like broadcasts and reductions.\n\nThe library embeds patterns similar to those in TensorFlow Federated but uses JAX's mapping and extension capabilities. It is tailored for scenarios involving large models and distributed training where both correctness and efficiency matter. The package depends on absl-py, jax, and jaxlib, and requires Python 3.9 or later. It is actively maintained and carries no known security vulnerabilities.","worth_installing":"Yes, if you are building distributed or federated machine learning systems in JAX and need MapReduce-style primitives with automatic differentiation. The package is actively maintained, carries no known vulnerabilities, and has low install friction. However, it remains early-stage (first release January 2025) with modest adoption; verify that its API and performance characteristics match your production requirements before committing to it."},"id":"drjax","links":{"html":"https://skillfed.io/packages/drjax","md":"https://skillfed.io/packages/drjax.md","pypi":"https://pypi.org/project/drjax/"},"maintenance":{"status":"active"},"meta":{"latest_release":"2026-06-15","license_spdx":null,"license_treatment":"permissive","name":"drjax","python_support":"supports_current","summary":"DrJAX - Scalable and Differentiable MapReduce Primitives in JAX."},"popularity":{"monthly_downloads":253355,"position":8521,"tier":"top_15000"},"security":{"n_vulnerabilities":0},"version":"0.2.0"}