Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors and unifies the Rotary Position Embedding (RoPE) implementations across multiple model architectures into a centralized, modular framework under lightx2v/common/ops/rope. It introduces a base RopeTemplate alongside specialized classes like TorchComplexRope, TorchRealRope, FlashInferRope, and ChunkedRope, as well as a builder utility to streamline configuration. While this significantly cleans up duplicate code, several critical issues were identified during the review: an uninitialized self.causal_torch_rope attribute in the audio transformer will cause an AttributeError at runtime; ChunkedRope.apply fails to independently handle asymmetric sequence lengths for q and k, which can leave k_out uninitialized; FlashInferRope.apply_single modifies the input tensor x inplace, causing unintended side effects; and the builder's use of issubclass on registered callables can trigger a TypeError if they are not actual classes.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| x_c = torch.view_as_complex(x.float().reshape(x.size(0), n, -1, 2)) | ||
| out = torch.view_as_real(x_c * pos_freqs.to(torch.complex64)).flatten(2) | ||
| return out.to(orig_dtype) | ||
| return self.causal_torch_rope.apply_single(x, pos_freqs.to(torch.complex64)).to(orig_dtype) |
There was a problem hiding this comment.
The attribute self.causal_torch_rope is used here but it is never initialized or defined in this class (or its parent classes). This will result in an AttributeError at runtime. Please initialize it in the constructor or define it as a module-level constant.
| if seq_len < q.shape[0]: | ||
| q_out[seq_len:], k_out[seq_len:] = q[seq_len:], k[seq_len:] |
There was a problem hiding this comment.
There is a correctness bug when q and k have different sequence lengths (e.g., during autoregressive decoding where q has length 1 and k has the full context length). If seq_len < q.shape[0] is False but seq_len < k.shape[0] is True, the remaining part of k_out will not be copied from k, leaving it uninitialized. They should be checked and copied independently.
| if seq_len < q.shape[0]: | |
| q_out[seq_len:], k_out[seq_len:] = q[seq_len:], k[seq_len:] | |
| if seq_len < q.shape[0]: | |
| q_out[seq_len:] = q[seq_len:] | |
| if seq_len < k.shape[0]: | |
| k_out[seq_len:] = k[seq_len:] |
| def apply_single(self, x: torch.Tensor, freqs, **kwargs): | ||
| output, _ = self.apply(x, x.clone(), freqs, **kwargs) | ||
| return output |
There was a problem hiding this comment.
In apply_single, x is passed directly as the first argument to self.apply, which calls _apply_eager and modifies it inplace. This will unexpectedly modify the input tensor x for the caller. To prevent this side effect, both arguments should be cloned.
| def apply_single(self, x: torch.Tensor, freqs, **kwargs): | |
| output, _ = self.apply(x, x.clone(), freqs, **kwargs) | |
| return output | |
| def apply_single(self, x: torch.Tensor, freqs, **kwargs): | |
| output, _ = self.apply(x.clone(), x.clone(), freqs, **kwargs) | |
| return output |
| if layout != "interleaved": | ||
| raise ValueError("TorchComplexRope only supports interleaved layout.") | ||
| return rope_class(compute_dtype=compute_dtype) | ||
| if issubclass(rope_class, RopeTemplate): |
There was a problem hiding this comment.
Using issubclass directly on rope_class will raise a TypeError if a non-class callable (like a factory function) is registered in ROPE_REGISTER. It is safer to check if it is a class first using isinstance(rope_class, type).
| if issubclass(rope_class, RopeTemplate): | |
| if isinstance(rope_class, type) and issubclass(rope_class, RopeTemplate): |
1 participant
No description provided.