Difficulty Configuring Nested Initializer Parameters

[BeeGass]

Context and Problem Description

I'm attempting to create a flexible configuration system for a machine learning project using hydra-zen, specifically for configuring observation layer initializers. My goal is to:

1. Set default initializers in the configuration
2. Allow runtime override of the initializer type
3. Allow runtime configuration of initializer parameters (e.g., standard deviation)

Specific Use Case

I want to be able to run configurations like:

python src/test.py \
    observation_layer.x_dim=32 \
    observation_layer.features=64 \
    observation_layer.kernel_init=normal \
    observation_layer.kernel_init.stddev=0.1

or override other intializers at the same time:

python src/test.py \
    observation_layer.x_dim=32 \
    observation_layer.features=64 \
    observation_layer.time_scale_init=normal \
    observation_layer.time_scale_init.stddev=0.5 \
    observation_layer.kernel_init=normal \ # override from lecun_normal
    observation_layer.kernel_init.stddev=0.1 

Current Implementation Attempt

from dataclasses import asdict
from dataclasses import dataclass
from typing import Literal
from typing import Optional

import jax.numpy as jnp
from flax.linen.initializers import normal
from flax.linen.initializers import uniform
from flax.linen.initializers import zeros
from flax.linen.initializers import zeros_init
from flax.nnx.nn.initializers import lecun_normal
from flax.nnx.nn.initializers import lecun_uniform
from flax.typing import Dtype
from flax.typing import Initializer
from hydra_zen import MISSING
from hydra_zen import builds
from hydra_zen import just
from hydra_zen import store
from hydra_zen.third_party.pydantic import pydantic_parser

from src.train.utils import custom_format_dict

default_kernel_init = lecun_normal()
default_bias_init = zeros_init()

ShardingRules = Optional[tuple[tuple[str, str], ...]]
TimeScaleType = Literal["scalar", "vector"]
ActivationType = Literal["exp", "tanh", "sigmoid", "relu"]
EinOptimumType = Literal[
    "auto",
    "auto-hq",
    "optimal",
    "dynamic-programming",
    "branch",
    "branch-all",
    "branch-2",
    "random-greedy-128",
    "random-greedy",
    "greedy",
]

@dataclass
class ObservationFirstOrderArgs:
    # Core dimensions
    x_dim: int = MISSING
    features: int = MISSING
    # Model configuration
    num_head: int = 1
    time_scale: TimeScaleType = "scalar"
    activation: ActivationType = "exp"
    # Parameter initialization
    kernel_init: Initializer = default_kernel_init
    time_scale_init: Initializer = normal(stddev=0.1)
    param_dtype: Dtype = jnp.float32
    dtype: Dtype = jnp.bfloat16
    # Sharding, optimization and RNG
    opt_path: EinOptimumType = "optimal"

# Store for kernel initializers
obs_kernel_init_store = store(group="observation_layer/kernel_init")
obs_kernel_init_store(builds(normal, stddev=0.1, populate_full_signature=True), name="normal")
obs_kernel_init_store(builds(lecun_normal, populate_full_signature=True), name="lecun_normal")
obs_kernel_init_store(builds(lecun_uniform, populate_full_signature=True), name="lecun_uniform")
obs_kernel_init_store(builds(uniform, scale=0.01, populate_full_signature=True), name="uniform")
obs_kernel_init_store(builds(zeros, populate_full_signature=True), name="zeros")

# Store for time scale initializers
obs_time_scale_init_store = store(group="observation_layer/time_scale_init")
obs_time_scale_init_store(builds(normal, stddev=0.1, populate_full_signature=True), name="normal")
obs_time_scale_init_store(builds(lecun_normal, populate_full_signature=True), name="lecun_normal")

# Store for dtypes
obs_param_dtype_store = store(group="observation_layer/param_dtype")
obs_param_dtype_store(just(jnp.float32), name="float32")
obs_param_dtype_store(just(jnp.bfloat16), name="bfloat16")

obs_dtype_store = store(group="observation_layer/dtype")
obs_dtype_store(just(jnp.float32), name="float32")
obs_dtype_store(just(jnp.bfloat16), name="bfloat16")

# Create the observation layer configs
ObservationFirstOrderConf = builds(
    ObservationFirstOrderArgs,
    hydra_defaults=[
        "_self_",
        {"kernel_init": "lecun_normal"},  # default kernel initializer
        {"time_scale_init": "normal"},  # default time scale initializer
        {"param_dtype": "float32"},  # default param dtype
        {"dtype": "bfloat16"},  # default dtype
    ],
    populate_full_signature=True,
    zen_dataclass={'kw_only': True},
    # Provide default values for required fields
    x_dim=MISSING,  # These will need to be specified when used
    features=MISSING,
    kernel_init=MISSING,  # Will be populated from defaults
    time_scale_init=MISSING,  # Will be populated from defaults
    param_dtype=MISSING,  # Will be populated from defaults
    dtype=MISSING  # Will be populated from defaults
)

# Store for observation layers
observation_layer_store = store(group="observation_layer")
observation_layer_store(ObservationFirstOrderConf, name="first_order")

# Main function
def main_fn(observation_layer=None):
    print(f"observation_layer configuration:\n{custom_format_dict(asdict(observation_layer))}\n\n")

store(
    main_fn,
    hydra_defaults=[
        "_self_",
        {"observation_layer": "first_order"}
    ]
)

if __name__ == "__main__":
    from hydra_zen import zen
    from hydra_zen.third_party.pydantic import validates_with_pydantic

    store.add_to_hydra_store()
    validated_main = validates_with_pydantic(main_fn, validator=pydantic_parser)

    zen(main_fn).hydra_main(
        config_name="main_fn",
        config_path=None,
        version_base="1.3",
    )

which works when I do the following:

uv run src/test.py \                                                                                                                                                                  
    observation_layer.x_dim=32 \
    observation_layer.features=64

output:

observation_layer configuration:
{
'x_dim': 32,
'features': 64,
'num_head': 1,
'time_scale': 'scalar',
'activation': 'exp',
'kernel_init': <function init from jax._src.nn.initializers>,
'time_scale_init': <function init from jax._src.nn.initializers>,
'param_dtype': <class 'jax.numpy.float32'>,
'dtype': <class 'jax.numpy.bfloat16'>,
'opt_path': 'optimal'
}

likewise when I do:

uv run src/test.py \
    observation_layer.x_dim=32 \
    observation_layer.features=64 \
    observation_layer.time_scale_init.stddev=0.5

output:

observation_layer configuration:
{
'x_dim': 32,
'features': 64,
'num_head': 1,
'time_scale': 'scalar',
'activation': 'exp',
'kernel_init': <function init from jax._src.nn.initializers>,
'time_scale_init': <function init from jax._src.nn.initializers>,
'param_dtype': <class 'jax.numpy.float32'>,
'dtype': <class 'jax.numpy.bfloat16'>,
'opt_path': 'optimal'
}

Current Limitations

When attempting to override initializer parameters, I encounter errors like:

command 1:

uv run src/test.py \
    observation_layer.x_dim=32 \
    observation_layer.features=64 \
    observation_layer.time_scale_init=normal \
    observation_layer.time_scale_init.stddev=0.5

output:

Could not override 'observation_layer.time_scale_init.stddev'.
To append to your config use +observation_layer.time_scale_init.stddev=0.5
Key 'stddev' is not in struct
    full_key: observation_layer.time_scale_init.stddev
    object_type=dict

command 2:

uv run src/test.py \
    observation_layer.x_dim=32 \
    observation_layer.features=64 \
    observation_layer.time_scale_init=normal \
    +observation_layer.time_scale_init.stddev=0.5

output (incorrectly stores the jax function):

observation_layer configuration:
{
'x_dim': 32,
'features': 64,
'num_head': 1,
'time_scale': 'scalar',
'activation': 'exp',
'kernel_init': <function init from jax._src.nn.initializers>,
'time_scale_init': {'stddev': 0.5},
'param_dtype': <class 'jax.numpy.float32'>,
'dtype': <class 'jax.numpy.bfloat16'>,
'opt_path': 'optimal'
}

command 3:

uv run src/test.py \
    observation_layer.x_dim=32 \
    observation_layer.features=64 \
    observation_layer.time_scale_init=normal \
    observation_layer/time_scale_init/stddev=0.5

output:

Could not override 'observation_layer/time_scale_init/stddev'.
To append to your config use +observation_layer/time_scale_init/stddev=0.5
Key 'observation_layer/time_scale_init/stddev' not in 'Builds_main_fn'
    full_key: observation_layer/time_scale_init/stddev
    object_type=Builds_main_fn

command 4:

uv run src/test.py \
    observation_layer.x_dim=32 \
    observation_layer.features=64 \
    observation_layer.time_scale_init=normal \
    +observation_layer/time_scale_init/stddev=0.5

output (again incorrectly/doesnt store the jax function):

observation_layer configuration:
{
'x_dim': 32,
'features': 64,
'num_head': 1,
'time_scale': 'scalar',
'activation': 'exp',
'kernel_init': <function init from jax._src.nn.initializers>,
'time_scale_init': 'normal',
'param_dtype': <class 'jax.numpy.float32'>,
'dtype': <class 'jax.numpy.bfloat16'>,
'opt_path': 'optimal'
}

Questions and Requests for Guidance

1. What is the recommended pattern for creating configurable initializers in hydra-zen?
2. How can I create a configuration that supports:
   • Default initializer selection
   • Runtime initializer type change
   • Runtime initializer parameter configuration
3. Are there any examples or patterns in the hydra-zen ecosystem that demonstrate this level of configuration flexibility?