Neural Network Config Manager

A configuration management system for neural network hyperparameters that integrates with JAX transformations.

Capabilities

Define network configuration dataclass

Create a dataclass that stores neural network hyperparameters and can be used seamlessly with JAX tree operations.

• A NetworkConfig dataclass with fields hidden_dim (int), num_layers (int), dropout_rate (float), and learning_rate (float) can be created and is compatible with JAX tree utilities @test
• The NetworkConfig dataclass can be flattened and unflattened using jax.tree_util.tree_flatten and jax.tree_util.tree_unflatten @test

Map transformations over configuration

Apply transformations to all numeric values in the configuration using JAX tree mapping.

• Given a NetworkConfig with values hidden_dim=128, num_layers=3, dropout_rate=0.1, learning_rate=0.001, applying jax.tree_map(lambda x: x * 2) doubles all numeric values @test

Support dictionary-style access

Enable dictionary-like access patterns for configuration fields.

• A MappableNetworkConfig dataclass can be accessed using dictionary syntax (e.g., config['hidden_dim']) to retrieve field values @test
• The MappableNetworkConfig supports iteration over keys using standard dictionary methods like .keys(), .values(), and .items() @test

Create configuration from tuple

Reconstruct a configuration dataclass from a tuple representation.

• A NetworkConfig instance can be created from a tuple (128, 3, 0.1, 0.001) using the from_tuple method, matching the order of fields @test

Implementation

@generates

API

"""Neural network configuration management."""

# Define your NetworkConfig dataclass here
# Define your MappableNetworkConfig dataclass here

Dependencies { .dependencies }

chex { .dependency }

Provides JAX-friendly dataclass decorators and testing utilities.