PennyLaneAI / PennyLaneAI/catalyst

Fix issues around frontend path discovery of externally built extension modules

Open
#2,730 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

housecleaning
Dominant language
Python
Stars
234
Forks
84
Avg merge
2d 15h
Merged PRs (30d)
66

Description

A few issues around path discovery of externally built extension modules:

1) Runtime nanobind module discovery repeatedly appends to path

There are two locations where this happens for the runtime registry module:

https://github.com/PennyLaneAI/catalyst/blob/c6f318051063316f2e23e4c01808744ed5c2c51f/frontend/catalyst/jax_primitives.py#L485

and

https://github.com/PennyLaneAI/catalyst/blob/c6f318051063316f2e23e4c01808744ed5c2c51f/frontend/catalyst/compiler.py#L92

Because these lines are run every-time functionality from the module is used, the system path would get polluted with more and more duplicate entries. This doesn't break anything but is not ideal.

A solution could perform this task once upon import for example, similar to how we handle the mlir_quantum module:
https://github.com/PennyLaneAI/catalyst/blob/c6f318051063316f2e23e4c01808744ed5c2c51f/frontend/catalyst/__init__.py#L55-L62

The thing to keep in mind is that these modules may not exist in some contexts, in particular for a CI docs build which uses only the Python code. So code importing such modules would fail in this context. Typically, we resolve this with mocking the modules in the docs config:

https://github.com/PennyLaneAI/catalyst/blob/c6f318051063316f2e23e4c01808744ed5c2c51f/doc/conf.py#L88-L90

2) The fix should consider a similar module that will be merged soon from the MLIR component to make sure all such modules are handled uniformly

See https://github.com/PennyLaneAI/catalyst/pull/2259 which introduced the new module.

3) The path mlir_quantum should consider custom build directories when not in INSTALLED mode.

See also https://github.com/PennyLaneAI/catalyst/pull/2713#discussion_r3126902923 for a discussion of this.

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 with frontend/catalyst/jax_primitives.py, frontend/catalyst/compiler.py, and frontend/catalyst/init.py, then review doc/conf.py and the linked module changes in pull request 2259. Check the discussion in pull request 2713 about custom build directories. Done means externally built modules are discovered uniformly without repeated path entries, while docs-only imports and non-installed custom builds continue to work.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
build-system
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.