(freqs_t, freqs_h, freqs_w)
| 930 | |
| 931 | # BroadCast and concatenate temporal and spaial frequencie (height and width) into a 3d tensor |
| 932 | def combine_time_height_width(freqs_t, freqs_h, freqs_w): |
| 933 | freqs_t = freqs_t[:, None, None, :].expand( |
| 934 | -1, grid_size_h, grid_size_w, -1 |
| 935 | ) # temporal_size, grid_size_h, grid_size_w, dim_t |
| 936 | freqs_h = freqs_h[None, :, None, :].expand( |
| 937 | temporal_size, -1, grid_size_w, -1 |
| 938 | ) # temporal_size, grid_size_h, grid_size_2, dim_h |
| 939 | freqs_w = freqs_w[None, None, :, :].expand( |
| 940 | temporal_size, grid_size_h, -1, -1 |
| 941 | ) # temporal_size, grid_size_h, grid_size_2, dim_w |
| 942 | |
| 943 | freqs = torch.cat( |
| 944 | [freqs_t, freqs_h, freqs_w], dim=-1 |
| 945 | ) # temporal_size, grid_size_h, grid_size_w, (dim_t + dim_h + dim_w) |
| 946 | freqs = freqs.view( |
| 947 | temporal_size * grid_size_h * grid_size_w, -1 |
| 948 | ) # (temporal_size * grid_size_h * grid_size_w), (dim_t + dim_h + dim_w) |
| 949 | return freqs |
| 950 | |
| 951 | t_cos, t_sin = freqs_t # both t_cos and t_sin has shape: temporal_size, dim_t |
| 952 | h_cos, h_sin = freqs_h # both h_cos and h_sin has shape: grid_size_h, dim_h |
no outgoing calls
no test coverage detected