🧅 Unions

ML configs often need mutually exclusive shapes: a backbone that is either a transformer or a conv net, an optimizer that is Adam or SGD, a schedule that is cosine or step. Protobuf's oneof encodes that exclusivity in the schema. py-gen-ml turns it into a typed Union on the generated Pydantic models.

📝 Defining a oneof

// oneof_demo.proto
syntax = "proto3";

package oneof_demo;

// Transformer configuration
message Transformer {
    // Number of layers
    uint32 num_layers = 1;
    // Number of heads
    uint32 num_heads = 2;
    // Activation function
    string activation = 3;
}

// Conv block
message ConvBlock {
    // Number of output channels
    uint32 out_channels = 1;
    // Kernel size
    uint32 kernel_size = 2;
    // Activation function
    string activation = 3;
}

// Convolutional neural network configuration
message ConvNet {
    // Conv layer configuration
    repeated ConvBlock layers = 1;
}

// Model configuration
message Model {
    oneof backbone {
        Transformer transformer = 1;
        ConvNet conv_net = 2;
    }
}

The generated code will look like this:

# Autogenerated code. DO NOT EDIT.
import typing
import py_gen_ml as pgml


class Transformer(pgml.YamlBaseModel):
    """Transformer configuration"""

    num_layers: int
    """Number of layers"""

    num_heads: int
    """Number of heads"""

    activation: str
    """Activation function"""


class ConvBlock(pgml.YamlBaseModel):
    """Conv block"""

    out_channels: int
    """Number of output channels"""

    kernel_size: int
    """Kernel size"""

    activation: str
    """Activation function"""


class ConvNet(pgml.YamlBaseModel):
    """Convolutional neural network configuration"""

    layers: typing.List[ConvBlock]
    """Conv layer configuration"""


class Model(pgml.YamlBaseModel):
    """Model configuration"""

    backbone: typing.Union[Transformer, ConvNet]

Notice backbone: typing.Union[Transformer, ConvNet]. Exactly one of the alternatives is expected at runtime.

📄 YAML shape

In YAML, nest the chosen variant under its field name. For a transformer backbone:

backbone:
  transformer:
    num_layers: 6
    num_heads: 8
    activation: gelu

Or for a conv net:

backbone:
  conv_net:
    layers:
      - out_channels: 64
        kernel_size: 3
        activation: relu

The generated JSON Schema validates that the payload matches one of the allowed variants.

🔧 Patches, sweeps, and typing

• Patches: each alternative can be patched independently. Unset fields on the patch leave the base alone.
• Sweeps: you can sweep fields inside the active variant the same way you sweep nested messages.
• Typing: your training code can branch with isinstance (or pattern matching) on the union members and stay type-checker friendly.

⚠️ Caveats

• A oneof is exclusive: don't set more than one alternative in the same config.
• Prefer oneof when the alternatives have different fields. If they share the same shape and only a label differs, an enum is usually simpler.