Skip to content

nos.common.types

nos.common.types

TensorSpec

Source code in nos/common/types.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
@dataclass(frozen=True)
class TensorSpec:
    shape: Optional[Tuple[Optional[int], ...]] = None
    """Base tensor specification with at most 4 dimensions.

    This class is used to capture the shape and dtype of a tensor.
    Tensor shapes are specified as a tuple of integers, where
    the first dimension is the batch size. Each of the dimensions
    are optional, and can be set to None to support dynamic dimension.
    For e.g., a model that supports variable batch size has
    shape=(None, 224, 224, 3).

    Examples:
        ImageSpec:
            - (H, W, C): (height, width, channels)
        EmbeddingSpec:
            - (D): (dims, )
    """
    dtype: str = None
    """Tensor dtype. (uint8, int32, int64, float32, float64)"""

    @field_validator("shape")
    @classmethod
    def validate_shape(cls, shape: Optional[Tuple[Optional[int], ...]]):
        """Validate the shape."""
        if shape and (len(shape) < 1 or len(shape) > 4):
            raise ValueError(f"Invalid tensor shape [shape={shape}].")
        else:
            return shape

    @field_validator("dtype", mode="before")
    @classmethod
    def validate_dtype(cls, dtype: str):
        """Validate the dtype."""
        if dtype and not hasattr(np, dtype):
            raise ValueError(f"Invalid dtype [dtype={dtype}].")
        else:
            return dtype

    @property
    def nbytes(self) -> Optional[int]:
        """Return the number of bytes required to store the tensor."""
        try:
            return np.prod(self.shape) * np.dtype(self.dtype).itemsize
        except TypeError:
            return None

shape class-attribute instance-attribute 

shape: Optional[Tuple[Optional[int], ...]] = None

Base tensor specification with at most 4 dimensions.

This class is used to capture the shape and dtype of a tensor. Tensor shapes are specified as a tuple of integers, where the first dimension is the batch size. Each of the dimensions are optional, and can be set to None to support dynamic dimension. For e.g., a model that supports variable batch size has shape=(None, 224, 224, 3).

Examples:

ImageSpec: - (H, W, C): (height, width, channels) EmbeddingSpec: - (D): (dims, )

dtype class-attribute instance-attribute 

dtype: str = None

Tensor dtype. (uint8, int32, int64, float32, float64)

nbytes property 

nbytes: Optional[int]

Return the number of bytes required to store the tensor.

validate_shape classmethod 

validate_shape(shape: Optional[Tuple[Optional[int], ...]])

Validate the shape.

Source code in nos/common/types.py 
39
40
41
42
43
44
45
46
@field_validator("shape")
@classmethod
def validate_shape(cls, shape: Optional[Tuple[Optional[int], ...]]):
    """Validate the shape."""
    if shape and (len(shape) < 1 or len(shape) > 4):
        raise ValueError(f"Invalid tensor shape [shape={shape}].")
    else:
        return shape

validate_dtype classmethod 

validate_dtype(dtype: str)

Validate the dtype.

Source code in nos/common/types.py 
48
49
50
51
52
53
54
55
@field_validator("dtype", mode="before")
@classmethod
def validate_dtype(cls, dtype: str):
    """Validate the dtype."""
    if dtype and not hasattr(np, dtype):
        raise ValueError(f"Invalid dtype [dtype={dtype}].")
    else:
        return dtype

ImageSpec

Image tensor specification with dimensions (H, W, C).

Source code in nos/common/types.py 
66
67
68
69
70
71
72
73
74
75
76
77
@dataclass(frozen=True)
class ImageSpec(TensorSpec):
    """Image tensor specification with dimensions (H, W, C)."""

    @field_validator("shape", mode="before")
    @classmethod
    def validate_shape(cls, shape: Tuple[Optional[int], ...]):
        """Validate the shape."""
        if shape and len(shape) != 3:
            raise ValidationError(f"Invalid image shape [shape={shape}].")
        else:
            return shape

validate_shape classmethod 

validate_shape(shape: Tuple[Optional[int], ...])

Validate the shape.

Source code in nos/common/types.py 
70
71
72
73
74
75
76
77
@field_validator("shape", mode="before")
@classmethod
def validate_shape(cls, shape: Tuple[Optional[int], ...]):
    """Validate the shape."""
    if shape and len(shape) != 3:
        raise ValidationError(f"Invalid image shape [shape={shape}].")
    else:
        return shape

EmbeddingSpec

Embedding tensor specification with dimensions (D).

Source code in nos/common/types.py 
80
81
82
83
84
85
86
87
88
89
90
91
@dataclass(frozen=True)
class EmbeddingSpec(TensorSpec):
    """Embedding tensor specification with dimensions (D)."""

    @field_validator("shape", mode="before")
    @classmethod
    def validate_shape(cls, shape: Tuple[Optional[int]]):
        """Validate the shape."""
        if shape and len(shape) != 1:
            raise ValidationError(f"Invalid embedding shape [shape={shape}].")
        else:
            return shape

validate_shape classmethod 

validate_shape(shape: Tuple[Optional[int]])

Validate the shape.

Source code in nos/common/types.py 
84
85
86
87
88
89
90
91
@field_validator("shape", mode="before")
@classmethod
def validate_shape(cls, shape: Tuple[Optional[int]]):
    """Validate the shape."""
    if shape and len(shape) != 1:
        raise ValidationError(f"Invalid embedding shape [shape={shape}].")
    else:
        return shape

Batch

Generic annotation/type-hint for batched data.

Inherits from typing.Annotated[T, x] (PEP 593) where T is the type, and x is the metadata. The metadata is tpyically ignored, but can be used to allow additional type checks and annotations on the type.

Source code in nos/common/types.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
class Batch(Generic[T]):
    """Generic annotation/type-hint for batched data.

    Inherits from typing.Annotated[T, x] (PEP 593) where T is the type,
    and x is the metadata. The metadata is tpyically ignored,
    but can be used to allow additional type checks and annotations
    on the type.
    """

    __slots__ = ()

    @typing._tp_cache
    def __class_getitem__(cls, params):
        """Support Batch[T, batch_size].

        Annotated requires atleast 2 parameters [type, metadata].
        Here `batch_size` is optional (i.e. Batch[T],
        is equivalent to Batch[T, None]).
        """
        if not isinstance(params, tuple):
            params = (params, None)
        if len(params) != 2:
            raise TypeError(f"Invalid Batch parameters (T, batch_size), provided params={params}.")
        object_type, batch_size = params
        if batch_size is not None:
            if isinstance(batch_size, int):
                if batch_size < 1 or batch_size >= 65536:
                    raise ValueError(f"Invalid batch size [batch_size={batch_size}].")
            else:
                raise TypeError(f"Invalid batch size type [type(batch_size)={type(batch_size)}].")
        return Annotated[cls, object_type, batch_size]

__class_getitem__ 

__class_getitem__(params)

Support Batch[T, batch_size].

Annotated requires atleast 2 parameters [type, metadata]. Here batch_size is optional (i.e. Batch[T], is equivalent to Batch[T, None]).

Source code in nos/common/types.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
@typing._tp_cache
def __class_getitem__(cls, params):
    """Support Batch[T, batch_size].

    Annotated requires atleast 2 parameters [type, metadata].
    Here `batch_size` is optional (i.e. Batch[T],
    is equivalent to Batch[T, None]).
    """
    if not isinstance(params, tuple):
        params = (params, None)
    if len(params) != 2:
        raise TypeError(f"Invalid Batch parameters (T, batch_size), provided params={params}.")
    object_type, batch_size = params
    if batch_size is not None:
        if isinstance(batch_size, int):
            if batch_size < 1 or batch_size >= 65536:
                raise ValueError(f"Invalid batch size [batch_size={batch_size}].")
        else:
            raise TypeError(f"Invalid batch size type [type(batch_size)={type(batch_size)}].")
    return Annotated[cls, object_type, batch_size]

TensorT

Source code in nos/common/types.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
class TensorT(Generic[T]):
    __slots__ = ()

    @typing._tp_cache
    def __class_getitem__(cls, params):
        """Support TensorT[type, tensor_spec].

        Annotated requires atleast 2 parameters [type, tensor_spec].
        Here `tensor_spec` is optional (i.e. TensorT[T],
        is equivalent to TensorT[T, None]).

        Examples:
            TensorT[np.ndarray, TensorSpec()] := Annotated[TensorT, np.ndarray, TensorSpec()]
            TensorT[torch.Tensor, TensorSpec()] := Annotated[TensorT, torch.Tensor, TensorSpec()]
        """
        if not isinstance(params, tuple):
            params = (params, TensorSpec())
        if len(params) != 2:
            raise TypeError(f"Invalid TensorT parameters (T, tensort_spec), provided params={params}.")
        object_type, tensor_spec = params
        if tensor_spec is not None:
            if not isinstance(tensor_spec, TensorSpec):
                raise TypeError(f"Invalid tensor_spec metadata [tensor_spec={type(tensor_spec)}].")
        return Annotated[cls, object_type, tensor_spec]

__class_getitem__ 

__class_getitem__(params)

Support TensorT[type, tensor_spec].

Annotated requires atleast 2 parameters [type, tensor_spec]. Here tensor_spec is optional (i.e. TensorT[T], is equivalent to TensorT[T, None]).

Examples:

TensorT[np.ndarray, TensorSpec()] := Annotated[TensorT, np.ndarray, TensorSpec()] TensorT[torch.Tensor, TensorSpec()] := Annotated[TensorT, torch.Tensor, TensorSpec()]

Source code in nos/common/types.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
@typing._tp_cache
def __class_getitem__(cls, params):
    """Support TensorT[type, tensor_spec].

    Annotated requires atleast 2 parameters [type, tensor_spec].
    Here `tensor_spec` is optional (i.e. TensorT[T],
    is equivalent to TensorT[T, None]).

    Examples:
        TensorT[np.ndarray, TensorSpec()] := Annotated[TensorT, np.ndarray, TensorSpec()]
        TensorT[torch.Tensor, TensorSpec()] := Annotated[TensorT, torch.Tensor, TensorSpec()]
    """
    if not isinstance(params, tuple):
        params = (params, TensorSpec())
    if len(params) != 2:
        raise TypeError(f"Invalid TensorT parameters (T, tensort_spec), provided params={params}.")
    object_type, tensor_spec = params
    if tensor_spec is not None:
        if not isinstance(tensor_spec, TensorSpec):
            raise TypeError(f"Invalid tensor_spec metadata [tensor_spec={type(tensor_spec)}].")
    return Annotated[cls, object_type, tensor_spec]

ImageT

Source code in nos/common/types.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
class ImageT(Generic[T]):
    __slots__ = ()

    @typing._tp_cache
    def __class_getitem__(cls, params):
        """Support TensorT[type, image_spec].

        Annotated requires atleast 2 parameters [type, image_spec].
        Here `image_spec` is optional (i.e. ImageT[T], is equivalent to TensorT[T, ImageSpec()]).

        Examples:
            ImageT[PIL.Image.Image, ImageSpec()] := Annotated[ImageT, Image, ImageSpec()]
        """
        if not isinstance(params, tuple):
            params = (params, ImageSpec())
        if len(params) != 2:
            raise TypeError(f"Invalid ImageT parameters (T, tensort_spec), provided params={params}.")
        object_type, image_spec = params
        if image_spec is not None:
            if not isinstance(image_spec, ImageSpec):
                raise TypeError(f"Invalid image_spec metadata [tensor_spec={type(image_spec)}].")
        return Annotated[cls, object_type, image_spec]

__class_getitem__ 

__class_getitem__(params)

Support TensorT[type, image_spec].

Annotated requires atleast 2 parameters [type, image_spec]. Here image_spec is optional (i.e. ImageT[T], is equivalent to TensorT[T, ImageSpec()]).

Examples:

ImageT[PIL.Image.Image, ImageSpec()] := Annotated[ImageT, Image, ImageSpec()]

Source code in nos/common/types.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
@typing._tp_cache
def __class_getitem__(cls, params):
    """Support TensorT[type, image_spec].

    Annotated requires atleast 2 parameters [type, image_spec].
    Here `image_spec` is optional (i.e. ImageT[T], is equivalent to TensorT[T, ImageSpec()]).

    Examples:
        ImageT[PIL.Image.Image, ImageSpec()] := Annotated[ImageT, Image, ImageSpec()]
    """
    if not isinstance(params, tuple):
        params = (params, ImageSpec())
    if len(params) != 2:
        raise TypeError(f"Invalid ImageT parameters (T, tensort_spec), provided params={params}.")
    object_type, image_spec = params
    if image_spec is not None:
        if not isinstance(image_spec, ImageSpec):
            raise TypeError(f"Invalid image_spec metadata [tensor_spec={type(image_spec)}].")
    return Annotated[cls, object_type, image_spec]