Skip to content

supaernova.configs.steps.backends

[docs] module supaernova.configs.steps.backends

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
from typing import Literal, get_args

from pydantic import ConfigDict

from .steps import StepConfig

TFBackend = Literal["tf", "tensorflow"]
TCHBackend = Literal["tch", "torch"]
Backend = TFBackend | TCHBackend

BACKENDS = {"TensorFlow": TFBackend, "PyTorch": TCHBackend}
BACKENDS_STR = ", ".join(
    f"{get_args(B)} for {backend}" for backend, B in BACKENDS.items()
)


class AbstractModelConfig(StepConfig):
    model_config: ConfigDict = ConfigDict(extra="allow")

    # === Required ===
    backend: Backend