[doc] feat: improve docstrings in torch_functional.py (#1345) #4730
Conversation
Standardized docstrings for core utility functions in torch_functional.py following Google-style documentation format: Functions improved: - gather_from_labels: Added complete docstring with example usage - logprobs_from_logits_flash_attn: Added docstring explaining Flash Attention usage - logprobs_from_logits_torch_npu: Added docstring for NPU implementation - logprobs_from_logits_naive: Added docstring for standard implementation - logprobs_from_logits_v2: Enhanced docstring explaining memory efficiency - clip_by_value: Added complete docstring with See Also reference - entropy_from_logits: Added docstring with mathematical formula - entropy_from_logits_with_chunking: Added docstring explaining chunking strategy - masked_sum: Added complete docstring with type hints - compute_grad_norm: Added docstring clarifying squared norm return - broadcast_dict_tensor: Added docstring with optimization note - allgather_dict_tensors: Added complete docstring with return type All improvements include: - Proper type hints on function signatures - Detailed Args and Returns sections - Notes for important implementation details - Mathematical formulas where applicable Contributes to volcengine#1345 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <[email protected]>
|
|
![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 significantly improves the documentation and type hints in verl/utils/torch_functional.py, making the code more readable and maintainable. The new Google-style docstrings are comprehensive and provide valuable context, examples, and notes on implementation details. The code changes are focused on documentation and do not alter functionality. I have one suggestion to improve a type hint for better accuracy and consistency with modern PyTorch practices.
| def logprobs_from_logits_v2(logits: torch.FloatTensor, labels): | ||
| """ | ||
| A memory efficient implementation of logprobs_from_logits | ||
| def logprobs_from_logits_v2(logits: torch.FloatTensor, labels: torch.Tensor) -> torch.Tensor: |
There was a problem hiding this comment.
The type hint torch.FloatTensor for logits is too restrictive. The function body explicitly handles other dtypes like bfloat16 in the else block. Using torch.Tensor would be more accurate and align with modern PyTorch practices, as torch.FloatTensor is a legacy alias for a torch.float32 tensor.
| def logprobs_from_logits_v2(logits: torch.FloatTensor, labels: torch.Tensor) -> torch.Tensor: | |
| def logprobs_from_logits_v2(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: |
vermouth1992
changed the title
3 participants
Summary
This PR improves the documentation quality of
verl/utils/torch_functional.pyby adding comprehensive Google-style docstrings to core utility functions.Functions Improved
gather_from_labelslogprobs_from_logits_flash_attnlogprobs_from_logits_torch_npulogprobs_from_logits_naivelogprobs_from_logits_v2clip_by_valueentropy_from_logitsentropy_from_logits_with_chunkingmasked_sumcompute_grad_normbroadcast_dict_tensorallgather_dict_tensorsImprovements Include
ArgsandReturnssectionsNotesections for important implementation detailsTest plan
Contributes to #1345
🤖 Generated with Claude Code