Skip to content

Add type hints and docstrings for ALiBi layer - #4472

Open
Ritabanm wants to merge 2 commits into
ml-explore:mainfrom
Ritabanm:main
Open

Add type hints and docstrings for ALiBi layer#4472
Ritabanm wants to merge 2 commits into
ml-explore:mainfrom
Ritabanm:main

Conversation

@Ritabanm

@Ritabanm Ritabanm commented Sep 7, 2026

Copy link
Copy Markdown

Add type hints and documentation for the ALiBi layer.

  • Added explicit Python parameter and return type annotations for create_alibi_matrix, create_alibi_slope, and the __call__ method in python/mlx/nn/layers/positional_encoding.py.
  • Added reStructuredText docstrings detailing parameters, return values, and references to the original paper.
  • Verified locally using python/tests/test_nn.py -k alibi.

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants