| 43 | from yateto.memory import DenseMemoryLayout |
| 44 | |
| 45 | class OptionalDimTensor(Tensor): |
| 46 | # dimSize = 1 is considered optional |
| 47 | def __init__(self, name, optName, optSize, optPos, shape, spp=None, memoryLayoutClass=DenseMemoryLayout, alignStride=False): |
| 48 | self._optName = optName |
| 49 | self._optSize = optSize |
| 50 | self._optPos = optPos |
| 51 | shape = self.insertOptDim(shape, (self._optSize,)) |
| 52 | super().__init__(name, shape, spp, memoryLayoutClass, alignStride) |
| 53 | |
| 54 | def hasOptDim(self): |
| 55 | return self._optSize > 1 |
| 56 | |
| 57 | def insertOptDim(self, sliceable, item): |
| 58 | if self.hasOptDim(): |
| 59 | return sliceable[0:self._optPos] + item + sliceable[self._optPos:] |
| 60 | return sliceable |
| 61 | |
| 62 | def __getitem__(self, indexNames): |
| 63 | indexNames = self.insertOptDim(indexNames, self._optName) |
| 64 | return IndexedTensor(self, indexNames) |
| 65 | |
| 66 | def optName(self): |
| 67 | return self._optName |
| 68 | |
| 69 | def optSize(self): |
| 70 | return self._optSize |
| 71 | |
| 72 | def optPos(self): |
| 73 | return self._optPos |
no outgoing calls
no test coverage detected