PennyLaneAI / PennyLaneAI/catalyst
[Dynamic shapes] Avoid using `_no_cleanup_dead_vars`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 234
- Forks
- 84
- Avg merge
- 2d 15h
- Merged PRs (30d)
- 66
Description
-
Expected behavior: We should avoid patching JAX in ways that are likely to not be supported in JAX. One of the instances we differ from JAX that we currently are investigating is a possible error in the JAX's function
clean_up_dead_varswhen using dynamic shapes. As a workaround we disableclean_up_dead_varsby overriding it and just returning None. -
Actual behavior: Disable of
clean_up_dead_vars.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by tracing the override named _no_cleanup_dead_vars and the use of JAX's clean_up_dead_vars in the dynamic-shapes path. Reproduce the reported behavior with a dynamic-shapes case, then determine whether the workaround can be removed while preserving correct compilation and cleanup behavior. The issue does not name files or tests, so locating the relevant entry points is part of the work.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100