Hardware Compatibility Validator

A utility for validating hardware requirements before running JAX computations. The validator should check device availability and ensure the runtime environment meets specified hardware requirements.

Capabilities

Validates GPU availability

• When check_gpu_available() is called and a GPU is available, it returns True without raising errors @test
• When check_gpu_available() is called and no GPU is available, it raises an AssertionError with an informative message @test

Validates TPU availability

• When check_tpu_available() is called and a TPU is available, it returns True without raising errors @test
• When check_tpu_available() is called and no TPU is available, it raises an AssertionError with an informative message @test

Validates specific device requirements

• When validate_devices(['gpu', 'cpu']) is called and all specified devices are available, it returns True @test
• When validate_devices(['gpu', 'tpu']) is called and some devices are unavailable, it raises an AssertionError indicating which devices are missing @test

Implementation

@generates

API

def check_gpu_available() -> bool:
    """
    Check if GPU hardware is available in the current environment.

    Returns:
        bool: True if GPU is available

    Raises:
        AssertionError: If no GPU is available
    """

def check_tpu_available() -> bool:
    """
    Check if TPU hardware is available in the current environment.

    Returns:
        bool: True if TPU is available

    Raises:
        AssertionError: If no TPU is available
    """

def validate_devices(device_list: list) -> bool:
    """
    Validate that all specified devices are available.

    Parameters:
        device_list: List of device type strings (e.g., ['gpu', 'cpu', 'tpu'])

    Returns:
        bool: True if all devices are available

    Raises:
        AssertionError: If any specified device is unavailable
    """

Dependencies { .dependencies }

chex { .dependency }

Provides hardware assertion utilities for JAX.

@satisfied-by