Distributed Array Validator

A utility for validating that arrays in distributed computations are properly sharded across devices.

Capabilities

Validates sharding of single arrays

• When given a sharded array, validates successfully without raising errors @test
• When given a non-sharded array, raises an AssertionError @test

Validates sharding of nested structures

• When given a nested dictionary with all arrays sharded, validates successfully @test
• When given a nested dictionary with at least one non-sharded array, raises an AssertionError @test

Validates sharding of collections

• When given a list of sharded arrays, validates successfully @test
• When given an empty structure, validates successfully @test

Implementation

@generates

API

def validate_sharding(tree_of_arrays):
    """
    Validates that all arrays in a PyTree structure are properly sharded across devices.

    Args:
        tree_of_arrays: A JAX array or a PyTree (nested structure) containing JAX arrays

    Returns:
        bool: True if all arrays are properly sharded

    Raises:
        AssertionError: If any array in the structure is not properly sharded
    """
    pass

Dependencies { .dependencies }

chex { .dependency }

Provides JAX testing utilities including array sharding validation.

@satisfied-by