Flattens a dict with array's as values into a 1D vector (or 2D with batch dimension)
| 71 | |
| 72 | |
| 73 | class FlattenTransform(AbstractTransform): |
| 74 | """ Flattens a dict with array's as values into a 1D vector (or 2D with batch dimension)""" |
| 75 | |
| 76 | def __init__(self, exp_type: str): |
| 77 | super().__init__(exp_type) |
| 78 | self._out_dim = sum(self.input_dim.values()) |
| 79 | |
| 80 | def transform(self, X: Dict[str, np.ndarray]) -> np.ndarray: |
| 81 | return np.concatenate([X[key].flatten() for key in INPUT_TYPES], axis=0) |
| 82 | # return np.concatenate([torch.flatten(subX) for subX in X.values()], dim=0) |
| 83 | |
| 84 | def batched_transform(self, X: Dict[str, np.ndarray]) -> np.ndarray: |
| 85 | return np.concatenate([X[key].reshape((X[key].shape[0], -1)) for key in INPUT_TYPES], axis=1) |
| 86 | # return torch.cat([torch.flatten(subX, start_dim=1).unsqueeze(1) for subX in X.values()], dim=1) |
| 87 | |
| 88 | |
| 89 | class RepeatGlobalsTransform(AbstractTransform): |
nothing calls this directly
no outgoing calls
no test coverage detected