First commit

pkgs/xformers/components/positional_embedding/base.py

@@ -0,0 +1,35 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from abc import ABCMeta, abstractmethod
+from dataclasses import asdict, dataclass
+from typing import Type, TypeVar
+
+import torch.nn as nn
+
+Self = TypeVar("Self", bound="PositionEmbedding")
+
+
+@dataclass
+class PositionEmbeddingConfig:
+    name: str
+    dim_model: int
+    seq_len: int
+
+
+class PositionEmbedding(nn.Module, metaclass=ABCMeta):
+    @abstractmethod
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__()
+
+    @classmethod
+    def from_config(cls: Type[Self], config: PositionEmbeddingConfig) -> Self:
+        # Generate the class inputs from the config
+        fields = asdict(config)
+
+        # Skip all Nones so that default values are used
+        fields = {k: v for k, v in fields.items() if v is not None}
+        return cls(**fields)