Unstack batch dimension and split into channels and alpha mask.
(x, batch_size, num_channels=3)
| 59 | |
| 60 | |
| 61 | def unstack_and_split(x, batch_size, num_channels=3): |
| 62 | """Unstack batch dimension and split into channels and alpha mask.""" |
| 63 | unstacked = einops.rearrange(x, '(b s) c h w -> b s c h w', b=batch_size) |
| 64 | channels, masks = torch.split(unstacked, [num_channels, 1], dim=2) |
| 65 | return channels, masks |
| 66 | |
| 67 | |
| 68 | class SlotAttention(nn.Module): |