r""" Flax Convolutional based multi-head attention block for diffusion-based VAE. Parameters: channels (:obj:`int`): Input channels num_head_channels (:obj:`int`, *optional*, defaults to `None`): Number of attention heads num_groups (:obj:`int
| 200 | |
| 201 | |
| 202 | class FlaxAttentionBlock(nn.Module): |
| 203 | r""" |
| 204 | Flax Convolutional based multi-head attention block for diffusion-based VAE. |
| 205 | |
| 206 | Parameters: |
| 207 | channels (:obj:`int`): |
| 208 | Input channels |
| 209 | num_head_channels (:obj:`int`, *optional*, defaults to `None`): |
| 210 | Number of attention heads |
| 211 | num_groups (:obj:`int`, *optional*, defaults to `32`): |
| 212 | The number of groups to use for group norm |
| 213 | dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32): |
| 214 | Parameters `dtype` |
| 215 | |
| 216 | """ |
| 217 | |
| 218 | channels: int |
| 219 | num_head_channels: int = None |
| 220 | num_groups: int = 32 |
| 221 | dtype: jnp.dtype = jnp.float32 |
| 222 | |
| 223 | def setup(self): |
| 224 | self.num_heads = self.channels // self.num_head_channels if self.num_head_channels is not None else 1 |
| 225 | |
| 226 | dense = partial(nn.Dense, self.channels, dtype=self.dtype) |
| 227 | |
| 228 | self.group_norm = nn.GroupNorm(num_groups=self.num_groups, epsilon=1e-6) |
| 229 | self.query, self.key, self.value = dense(), dense(), dense() |
| 230 | self.proj_attn = dense() |
| 231 | |
| 232 | def transpose_for_scores(self, projection): |
| 233 | new_projection_shape = projection.shape[:-1] + (self.num_heads, -1) |
| 234 | # move heads to 2nd position (B, T, H * D) -> (B, T, H, D) |
| 235 | new_projection = projection.reshape(new_projection_shape) |
| 236 | # (B, T, H, D) -> (B, H, T, D) |
| 237 | new_projection = jnp.transpose(new_projection, (0, 2, 1, 3)) |
| 238 | return new_projection |
| 239 | |
| 240 | def __call__(self, hidden_states): |
| 241 | residual = hidden_states |
| 242 | batch, height, width, channels = hidden_states.shape |
| 243 | |
| 244 | hidden_states = self.group_norm(hidden_states) |
| 245 | |
| 246 | hidden_states = hidden_states.reshape((batch, height * width, channels)) |
| 247 | |
| 248 | query = self.query(hidden_states) |
| 249 | key = self.key(hidden_states) |
| 250 | value = self.value(hidden_states) |
| 251 | |
| 252 | # transpose |
| 253 | query = self.transpose_for_scores(query) |
| 254 | key = self.transpose_for_scores(key) |
| 255 | value = self.transpose_for_scores(value) |
| 256 | |
| 257 | # compute attentions |
| 258 | scale = 1 / math.sqrt(math.sqrt(self.channels / self.num_heads)) |
| 259 | attn_weights = jnp.einsum("...qc,...kc->...qk", query * scale, key * scale) |