conditional

conditional

ConditionalEncoder

ConditionalEncoder(input_specification, pretrained=False, frozen=False, device='cuda' if is_available() else 'cpu', model_dtype=None)

Bases: EncodingMixin

Encoder that conditions its outputs based on an external feature.

Subclasses implement encode(). The base class forward() handles temporal flatten/unflatten — encode() receives tensors without a time dimension.

Initialize conditional encoder.

Parameters:

Name	Type	Description	Default
input_specification	EncoderInput	Structured input specification with conditioning key.	required
pretrained	bool	Whether to use pretrained weights.	False
frozen	bool	Whether to freeze encoder weights.	False
device	str | None	Device to place the encoder on.	'cuda' if is_available() else 'cpu'
model_dtype	str | None	Precision string from experiment config (e.g. "bf16-mixed").	None

Source code in src/versatil/models/encoding/encoders/conditional.py

def __init__(
    self,
    input_specification: EncoderInput,
    pretrained: bool = False,
    frozen: bool = False,
    device: str | None = "cuda" if torch.cuda.is_available() else "cpu",
    model_dtype: str | None = None,
):
    """Initialize conditional encoder.

    Args:
        input_specification: Structured input specification with conditioning key.
        pretrained: Whether to use pretrained weights.
        frozen: Whether to freeze encoder weights.
        device: Device to place the encoder on.
        model_dtype: Precision string from experiment config (e.g. ``"bf16-mixed"``).
    """
    if not input_specification.conditioning_key:
        raise ValueError("Conditional encoder requires conditioning_key")
    super().__init__(
        input_specification=input_specification,
        pretrained=pretrained,
        frozen=frozen,
        device=device,
        model_dtype=model_dtype,
    )
    self.condition_key = input_specification.conditioning_key

forward

forward(inputs, conditioning)

Forward pass with temporal flatten/unflatten.

Parameters:

Name	Type	Description	Default
inputs	dict[str, Tensor]	Dict mapping input_keys to tensors with temporal dimension.	required
conditioning	Tensor	Conditioning tensor from another encoder.	required

Returns:

Type	Description
dict[str, Tensor]	Dict mapping feature names to feature tensors with temporal dimension.

Source code in src/versatil/models/encoding/encoders/conditional.py

def forward(
    self,
    inputs: dict[str, torch.Tensor],
    conditioning: torch.Tensor,
) -> dict[str, torch.Tensor]:
    """Forward pass with temporal flatten/unflatten.

    Args:
        inputs: Dict mapping input_keys to tensors with temporal dimension.
        conditioning: Conditioning tensor from another encoder.

    Returns:
        Dict mapping feature names to feature tensors with temporal dimension.
    """
    self._validate_inputs(inputs)
    inputs, conditioning, batch_size, temporal_length = self._flatten_temporal(
        inputs, conditioning
    )
    outputs = self.encode(inputs, conditioning)
    return self._unflatten_temporal(outputs, batch_size, temporal_length)

encode abstractmethod

encode(inputs, conditioning)

Encode inputs with conditioning, without temporal dimension.

Parameters:

Name	Type	Description	Default
inputs	dict[str, Tensor]	Dict mapping input_keys to tensors without temporal dimension.	required
conditioning	Tensor	Conditioning tensor (B, D).	required

Returns:

Type	Description
dict[str, Tensor]	Dict mapping feature names to feature tensors.

Source code in src/versatil/models/encoding/encoders/conditional.py

@abstractmethod
def encode(
    self,
    inputs: dict[str, torch.Tensor],
    conditioning: torch.Tensor,
) -> dict[str, torch.Tensor]:
    """Encode inputs with conditioning, without temporal dimension.

    Args:
        inputs: Dict mapping input_keys to tensors without temporal dimension.
        conditioning: Conditioning tensor (B, D).

    Returns:
        Dict mapping feature names to feature tensors.
    """
    raise NotImplementedError