How to properly compose configs with uninstantiated class objects and their arguments in a hierarchical structure?

[BeeGass]

The Problem

I'm working with a modular deep learning codebase using flax/NNX where I have a generic RNN model that can accept different types of RNN cells. I need to be able to configure both:

1. The specific RNN cell class to use (which should remain uninstantiated)
2. The arguments for that cell class

While maintaining the ability to easily swap between different cell types via hydra's configuration system.

Here's a simplified version demonstrating the core challenge:

# Base cell interface
class RNNCell(nnx.Module):
    def initialize_carry(self, input_shape: tuple[int, ...], rngs: Optional[nnx.Rngs] = None) -> Any:
        raise NotImplementedError()
        
    def __call__(self, carry: Any, inputs: jax.Array) -> tuple[Any, jax.Array]:
        raise NotImplementedError()

class SimpleRNNCell(RNNCell):
    def __init__(self, in_features: int, hidden_features: int, rngs: nnx.Rngs):
        ...

class LSTMCell(RNNCell):
    def __init__(self, in_features: int, hidden_features: int, rngs: nnx.Rngs):
        ...

# The RNN module that needs to be configured
class RNN(nnx.Module):
    def __init__(
        self,
        cell_type: Type[RNNCell],  # Should be uninstantiated class
        cell_kwargs: dict[str, Any] | None = None,
        *,
        rngs: Optional[nnx.Rngs] = None,
    ):
        self.cell = cell_type(rngs=rngs, **cell_kwargs)
        ...

My attempt at configuration:

from dataclasses import MISSING
from hydra_zen import builds, hydrated_dataclass, instantiate, store, zen


@hydrated_dataclass(target=dict)
class SimpleRNNCellArgs:
    in_features: int = MISSING
    hidden_features: int = MISSING


@hydrated_dataclass(target=dict)
class LSTMCellArgs:
    in_features: int = MISSING
    hidden_features: int = MISSING


# Store cell configs
store(SimpleRNNCellArgs, group="cell_type_args", name="simple_rnn_args")
store(LSTMCellArgs, group="cell_type_args", name="lstm_args")

store(SimpleRNNCell, group="cell", name="simple_rnn")
store(LSTMCell, group="cell", name="lstm")


# Define RNN config that should compose with a cell config
@hydrated_dataclass(target=dict)
class RNNArgs:
    """Config for RNN that should compose with a cell"""

    cell_type: Union[SimpleRNNCell, LSTMCell] = MISSING
    cell_kwargs: Union[SimpleRNNCellArgs, LSTMCellArgs] = MISSING


store(
    RNNArgs,
    group="models",
    name="rnn",
    hydra_defaults=["_self_", dict(cell_type="simple_rnn"), dict(cell_kwargs="simple_rnn_args")],
)


def task(model: RNNArgs):
    """Task showing challenge with instantiation"""
    config = instantiate(model)
    print(f"Model Args: {config}")

    # The challenge: How do we properly compose these so that:
    # 1. The cell_type remains uninstantiated until needed
    # 2. The cell kwargs are properly passed through
    cell = config["cell_type"]
    cell_kwargs = config["cell_kwargs"]

    # Goal: Be able to do something like:
    # rnn = RNN(cell_type=cell, cell_kwargs=cell_kwargs)


TaskConfig = builds(
    task, populate_full_signature=True, hydra_defaults=["_self_", dict(models="rnn")]
)

store(TaskConfig, name="my_task")

if __name__ == "__main__":
    store.add_to_hydra_store()
    zen(task).hydra_main(version_base="1.3", config_name="my_task", config_path=None)

The Challenge

The key challenges I'm running into are:

1. How to properly store and reference the uninstantiated cell classes (SimpleRNNCell, LSTMCell) in the config system
2. How to properly compose the cell class with its arguments in a way that:
   • Keeps the cell class uninstantiated until needed
   • Correctly passes through the cell's configuration parameters
   • Allows easy swapping between different cell types via hydra's config system
3. How to structure the configs to maintain type safety and configuration validation

I'm unsure how to properly:

1. Keep the cell_type as an uninstantiated class while properly passing its arguments
2. Structure the configs to make cell types cleanly swappable
3. Maintain proper typing throughout

What's the recommended pattern in hydra-zen for this type of hierarchical config where we need to compose uninstantiated classes with their arguments?

Thanks in advance. This has been boggling my mind so any help is appreciated.