Skip to content

Clarify forward() type signatures and add docstrings#215

Open
LiudengZhang wants to merge 1 commit into
prescient-design:mainfrom
LiudengZhang:fix/forward-type-signature
Open

Clarify forward() type signatures and add docstrings#215
LiudengZhang wants to merge 1 commit into
prescient-design:mainfrom
LiudengZhang:fix/forward-type-signature

Conversation

@LiudengZhang

Copy link
Copy Markdown

References and relevant issues

Closes #172

Summary

  • Add NumPy-style docstrings with parameter descriptions and shape annotations across 8 model files
  • Fix incorrect return type annotations (e.g., forward() returning a dict but annotated as Tensor)
  • No logic changes — annotations and docstrings only

Test plan

  • ruff check passes
  • Existing tests pass (no behavior changes)

🤖 Generated with Claude Code

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Confusing forward Type Signature

1 participant