drjax

import jax import jax.numpy as jnp from drjax import map_reduce # Define the 'map' function: square an item def mapper_fn(item): return item ** 2 # Define the 'reduce' function: sum two items def reducer_fn(a, b): return a + b # Generate some input data across 'shards' # For a single device, this is just a batch dimension num_shards = 4 num_items_per_shard = 10 data = jnp.arange(num_shards * num_items_per_shard).reshape(num_shards, num_items_per_shard) # Use map_reduce to apply mapper_fn to each item, then reducer_fn across results result = map_reduce( mapper_fn=mapper_fn, reducer_fn=reducer_fn, inputs=data ) print(f"Input data shape: {data.shape}") print(f"Input data (first shard): {data[0]}") print(f"Result: {result}") # Expected result: jnp.sum(jnp.arange(40)**2) == 21090