| 501 | |
| 502 | @dataclass |
| 503 | class LazyTensor: |
| 504 | _load: Callable[[], Tensor] |
| 505 | shape: list[int] |
| 506 | data_type: DataType |
| 507 | description: str |
| 508 | |
| 509 | def load(self) -> Tensor: |
| 510 | ret = self._load() |
| 511 | # Should be okay if it maps to the same numpy type? |
| 512 | assert ret.data_type == self.data_type or (self.data_type.dtype == ret.data_type.dtype), \ |
| 513 | (self.data_type, ret.data_type, self.description) |
| 514 | return ret |
| 515 | |
| 516 | def astype(self, data_type: DataType) -> LazyTensor: |
| 517 | self.validate_conversion_to(data_type) |
| 518 | |
| 519 | def load() -> Tensor: |
| 520 | return self.load().astype(data_type) |
| 521 | return LazyTensor(load, self.shape, data_type, f'convert({data_type}) {self.description}') |
| 522 | |
| 523 | def validate_conversion_to(self, data_type: DataType) -> None: |
| 524 | if data_type != self.data_type and data_type.name not in self.data_type.valid_conversions: |
| 525 | raise ValueError(f'Cannot validate conversion from {self.data_type} to {data_type}.') |
| 526 | |
| 527 | |
| 528 | LazyModel: TypeAlias = 'dict[str, LazyTensor]' |
no outgoing calls
no test coverage detected