pymc-devs / pymc-devs/pytensor

Module-level Op instances break LSP signatures, Sphinx docs, and source links for expm, det, inv, lstsq

Open
#2,184 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
644
Forks
208
Avg merge
2d 14h
Merged PRs (30d)
16

Description

Problem

Several public API symbols in pytensor.tensor.linalg are bare Op instances (or Blockwise(Op()) instances), not functions with real signatures. This causes three issues:

  1. LSPs (pyright, mypy, jedi) can only show the generic Op.__call__ signature (*inputs, name=None, return_list=False, **kwargs) — no parameter names, types, or docstrings.
  2. Sphinx autodoc renders the same generic signature and pulls the docstring from Blockwise instead of the actual operation.
  3. "Source" links in the docs point to pytensor/tensor/blockwise.py rather than the operation's actual home (e.g. products.py, summary.py, inverse.py).

Affected symbols:

Symbol File Current definition Docs source link points to
expm products.py:81 expm = Blockwise(Expm()) blockwise.py
det summary.py:81 det = Blockwise(Det()) blockwise.py
inv, matrix_inverse inverse.py:166 inv = matrix_inverse = Blockwise(MatrixInverse()) blockwise.py
lstsq solvers/lstsq.py:36 lstsq = Lstsq() solvers/lstsq.py

Documentation examples:

Compare with properly wrapped functions like solve, cholesky, svd, norm, tensorsolve which all show their actual parameter signatures, domain-specific docstrings, and correct source links.

Possible approaches

A) Function wrapper — The codebase already uses this pattern in tensor/math.py:

_matmul = Blockwise(_dot, name="Matmul")      # private Op instance

def matmul(x1, x2, dtype=None):                 # public function with real signature
    """Compute the matrix product of two tensor variables.
    ...
    """
    return _matmul(x1, x2)

For each affected symbol, this would mean keeping the Op instance private and exposing a wrapper function with proper annotations and docstring.

B) .pyi stub files — Add products.pyi, summary.pyi, etc. with typed signatures. Helps LSPs but doesn't fix Sphinx autodoc, and adds maintenance burden keeping stubs in sync.

C) typing.cast with a Callable annotation — Helps LSPs but doesn't fix Sphinx.

D) Instance __doc__ — e.g. a docstring after the assignment. Fixes the rendered description in Sphinx but not the signature, source link, or LSP behavior.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by comparing the public wrappers in tensor/math.py with the affected definitions in products.py, summary.py, inverse.py, and solvers/lstsq.py. Inspect how the linalg API is rendered by Sphinx and how its source links are generated. Done means the affected symbols expose operation-specific signatures and docstrings to LSPs and Sphinx, with source links pointing to their defining files.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
api, developer-experience, documentation
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.