PennyLaneAI / PennyLaneAI/catalyst

Function for setting and resetting global changes during compilation

Open
#913 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

There are several places where we modify JAX during compilation.

Just to list some:

# Required for JAX tracer objects as PennyLane wires.
# pylint: disable=unnecessary-lambda
setattr(jax.interpreters.partial_eval.DynamicJaxprTracer, "__hash__", lambda x: id(x))

# This flag cannot be set in ``QJIT.get_mlir()`` because values created before
# that function is called must be consistent with the JAX configuration value.
jax.config.update("jax_enable_x64", True)

Patchers... (see jax_extras, see jax_transient_config)

And we also have a global context to see whether or not we are running or jax via the EvaluationContext.

With callbacks, this now changes the assumption that if we are tracing, we will never go back to the python environment. We should have a function that is able to save the configuration before we trace, change however we want it, reset it during callbacks, and reset it back to what we need once we exit the callback scope.

Note, could we instead of changing jax.interpreters.partial_eval.DynamicJaxprTracer and adding a hash, can't we change pennylane wire utilities to find whether the wire is jax.interpreters.partial_eval.DynamicJaxprTracer and compute the id as its hash instead of modifying jax itself?

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 reviewing the existing jax_extras, jax_transient_config, QJIT.get_mlir(), and EvaluationContext references described in the issue. Trace how compilation and callbacks currently modify or observe global state, then define what must be saved, changed, and restored, including whether the DynamicJaxprTracer hash workaround belongs in PennyLane wire utilities instead.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.