The expected shape and dtype of a tensor
| 44 | |
| 45 | @dataclasses.dataclass |
| 46 | class TensorSpec: |
| 47 | """The expected shape and dtype of a tensor""" |
| 48 | shape: Tuple[int] |
| 49 | dtype: np.dtype |
| 50 | |
| 51 | def __post_init__(self): |
| 52 | assert all(x >= 0 for x in self.shape if x is not None) |
| 53 | |
| 54 | def extend(self, n, dim=0): |
| 55 | shape = list(self.shape) |
| 56 | shape[dim] += n |
| 57 | return TensorSpec(tuple(shape), self.dtype) |
| 58 | |
| 59 | def __mul__(self, other): |
| 60 | assert isinstance(other, int) |
| 61 | return TensorSpec(tuple(s*other for s in self.shape), dtype=self.dtype) |
| 62 | |
| 63 | @classmethod |
| 64 | def build(cls, data): |
| 65 | if data is None: |
| 66 | return None |
| 67 | return TensorSpec(data.shape, data.dtype) |
| 68 | |
| 69 | @classmethod |
| 70 | def get_spec(cls, src: TokenizedVisionData) -> Dict[str, 'TensorSpec']: |
| 71 | spec = {} if src.other_data is None else src.other_data |
| 72 | for k in ["tokens", "images", "image_masks", "token_pooling", "low_res_token_pooling"]: |
| 73 | v = getattr(src, k) |
| 74 | if v is not None: |
| 75 | spec[k] = TensorSpec.build(getattr(src, k)) |
| 76 | return spec |
| 77 | |
| 78 | @staticmethod |
| 79 | def max(*other: 'TensorSpec'): |
| 80 | other = [o for o in other if o is not None] |
| 81 | if not other: |
| 82 | return None |
| 83 | rank = len(other[0].shape) |
| 84 | dtype = other[0].dtype |
| 85 | assert all(rank == len(o.shape) for o in other) |
| 86 | assert all(dtype == o.dtype for o in other) |
| 87 | return TensorSpec( |
| 88 | [max([o.shape[i] for o in other]) for i in range(rank)], |
| 89 | dtype |
| 90 | ) |
| 91 | |
| 92 | def max_dictionaries(*other: Dict[str, 'TensorSpec']): |
| 93 | return {k: TensorSpec.max(*(o[k] for o in other if k in o)) for k in get_all_keys(other)} |
| 94 | |
| 95 | |
| 96 | @dataclasses.dataclass |
no outgoing calls
no test coverage detected