(self, ids: torch.Tensor)
| 697 | self.axes_dim = axes_dim |
| 698 | |
| 699 | def forward(self, ids: torch.Tensor) -> torch.Tensor: |
| 700 | n_axes = ids.shape[-1] |
| 701 | cos_out = [] |
| 702 | sin_out = [] |
| 703 | pos = ids.float() |
| 704 | is_mps = ids.device.type == "mps" |
| 705 | freqs_dtype = torch.float32 if is_mps else torch.float64 |
| 706 | for i in range(n_axes): |
| 707 | cos, sin = get_1d_rotary_pos_embed( |
| 708 | self.axes_dim[i], pos[:, i], repeat_interleave_real=True, use_real=True, freqs_dtype=freqs_dtype |
| 709 | ) |
| 710 | cos_out.append(cos) |
| 711 | sin_out.append(sin) |
| 712 | freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) |
| 713 | freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) |
| 714 | return freqs_cos, freqs_sin |
| 715 | |
| 716 | |
| 717 | class TimestepEmbedding(nn.Module): |
nothing calls this directly
no test coverage detected