| 239 | |
| 240 | |
| 241 | def broadcat(tensors, dim=-1): |
| 242 | num_tensors = len(tensors) |
| 243 | shape_lens = set(list(map(lambda t: len(t.shape), tensors))) |
| 244 | assert len(shape_lens) == 1, "tensors must all have the same number of dimensions" |
| 245 | shape_len = list(shape_lens)[0] |
| 246 | dim = (dim + shape_len) if dim < 0 else dim |
| 247 | dims = list(zip(*map(lambda t: list(t.shape), tensors))) |
| 248 | expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim] |
| 249 | assert all( |
| 250 | [*map(lambda t: len(set(t[1])) <= 2, expandable_dims)] |
| 251 | ), "invalid dimensions for broadcastable concatentation" |
| 252 | max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims)) |
| 253 | expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims)) |
| 254 | expanded_dims.insert(dim, (dim, dims[dim])) |
| 255 | expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims))) |
| 256 | tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes))) |
| 257 | return torch.cat(tensors, dim=dim) |
| 258 | |
| 259 | |
| 260 | def rotate_half(x): |