def read_dynamic(hf: h5py.File, key: str, begin: int, end: int, *, shared: bool = False) -> TensorList: try: offsets_ds = hf[f"{key}_offsets"] data_ds = hf[f"{key}_data"] except LookupError: return TensorList.empty(num_tensors=end - begin) allocator = allocate_shared_tensor if shared else torch.empty offsets = allocator((end - begin + 1, ), dtype=torch.long) offsets_ds.read_direct(offsets.numpy(), source_sel=np.s_[begin:end + 1]) data_begin = offsets[0].item() data_end = offsets[-1].item() data = allocator((data_end - data_begin, ), dtype=torch.long) # Needed because https://github.com/h5py/h5py/issues/870. if data_end - data_begin > 0: data_ds.read_direct(data.numpy(), source_sel=np.s_[data_begin:data_end]) offsets -= int(offsets[0]) return TensorList(offsets, data)
def test_from_tensor(self): self.assertEqual( EntityList.from_tensor(torch.tensor([3, 4], dtype=torch.long)), EntityList( torch.tensor([3, 4], dtype=torch.long), TensorList.empty(num_tensors=2) ), )
def test_empty(self): self.assertEqual( EntityList.empty(), EntityList(torch.empty((0, ), dtype=torch.long), TensorList.empty()), )
def from_tensor(cls, tensor: LongTensorType) -> 'EntityList': if tensor.dim() != 1: raise ValueError("Expected 1D tensor, got %dD" % tensor.dim()) tensor_list = TensorList.empty(num_tensors=tensor.shape[0]) return cls(tensor, tensor_list)
def empty(cls) -> 'EntityList': return cls(torch.empty((0, ), dtype=torch.long), TensorList.empty())