Neural Network Parameter Validator

A utility for validating neural network parameters represented as PyTree structures. The validator ensures that parameter trees contain valid numerical values and can compare parameter states across training checkpoints.

Capabilities

Validate finite values in parameter trees

• Given a PyTree containing arrays with finite values (no NaN or infinity), the validator confirms all values are finite @test
• Given a PyTree containing arrays with NaN values, the validator detects and reports the invalid values @test
• Given a PyTree containing arrays with infinity values, the validator detects and reports the invalid values @test

Compare parameter trees for equality

• Given two PyTrees with identical structure and values, the validator confirms they are equal @test
• Given two PyTrees with identical structure but different values, the validator detects the difference @test
• Given two PyTrees with different structures, the validator detects the structural difference @test

Implementation

@generates

API

def validate_parameters_finite(params):
    """
    Validates that all numerical values in a parameter PyTree are finite.

    Args:
        params: A PyTree (nested dict/list/tuple) containing JAX arrays

    Raises:
        AssertionError: If any non-finite values (NaN or infinity) are found
    """
    pass

def compare_parameter_trees(params1, params2):
    """
    Compares two parameter PyTrees for exact equality in structure and values.

    Args:
        params1: First PyTree containing JAX arrays
        params2: Second PyTree containing JAX arrays

    Raises:
        AssertionError: If the trees differ in structure or values
    """
    pass

Dependencies { .dependencies }

chex { .dependency }

Provides JAX testing and assertion utilities including tree validation functions.