mllam/neural-lam

Add type hints to methods in interaction_net.py

Open

#448 opened on Mar 20, 2026

View on GitHub
 (0 comments) (0 reactions) (0 assignees)Python (276 forks)auto 404
good first issuemaintenance

Repository metrics

Stars
 (282 stars)
PR merge metrics
 (PR metrics pending)

Description

Summary

The methods in neural_lam/interaction_net.py currently lack type annotations. mypy is already configured in pre-commit but cannot fully validate this module without proper type hints.

Methods affected

  • InteractionNet.forward
  • InteractionNet.message
  • InteractionNet.aggregate
  • SplitMLPs.forward

Proposed Fix

Add torch.Tensor, Optional[torch.Tensor], Tuple and Union type hints to all method signatures and return types.

Why This Matters

  • mypy is already set up in .pre-commit-config.yaml
  • Type hints improve contributor experience and catch bugs early
  • Consistent with type hints already added in metrics.py (see #446)

Contributor guide

Research direction
Review the type hints added in metrics.py (PR #446) for reference. Check the mypy settings in .pre commit config.yaml. Look at the signatures of the methods listed and add type annotations for arguments and return types.
Tech stack
pythonpytorch
Domain
machine learning
Issue type
Refactor
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
2
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
1-3 hours
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
Active
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
Clear
Prerequisites
GitPython
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
85