[BUG] abstracted_axes throws AttributeError in Pennylane QJIT

[sengthai]

Expected behavior

The code should run without errors and produce 1 and 4.

Actual behavior

Throws the error message:


AttributeError: 'tensor' object has no attribute 'aval'. Did you mean: 'ravel'?

Additional information

No response

Source code

import pennylane as qml
from pennylane import qjit
import pennylane.numpy as jnp

dev = qml.device("lightning.qubit", wires=1)

@qjit(verbose=True, abstracted_axes={0:"n"})
def sum_abstracted(arr):
    return jnp.sum(arr)

print(sum_abstracted(jnp.array([1])))  
print(sum_abstracted(jnp.array([1, 1, 2])))

Tracebacks

Traceback (most recent call last):
  File "/Users/sengthai.heng/Developer/test/pennylane_qjit.py", line 15, in <module>
    print(sum_abstracted(jnp.array([1])))  
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sengthai.heng/Developer/test/.venv/lib/python3.12/site-packages/pennylane/logging/decorators.py", line 61, in wrapper_entry
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sengthai.heng/Developer/test/.venv/lib/python3.12/site-packages/catalyst/jit.py", line 606, in __call__
    return self.run(args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sengthai.heng/Developer/test/.venv/lib/python3.12/site-packages/catalyst/debug/instruments.py", line 145, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/Users/sengthai.heng/Developer/test/.venv/lib/python3.12/site-packages/pennylane/logging/decorators.py", line 61, in wrapper_entry
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sengthai.heng/Developer/test/.venv/lib/python3.12/site-packages/catalyst/jit.py", line 820, in run
    results = self.compiled_function(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sengthai.heng/Developer/test/.venv/lib/python3.12/site-packages/catalyst/compiled_functions.py", line 337, in __call__
    dynamic_args = get_implicit_and_explicit_flat_args(
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sengthai.heng/Developer/test/.venv/lib/python3.12/site-packages/catalyst/jax_extras/tracing.py", line 491, in get_implicit_and_explicit_flat_args
    implicit_args = _extract_implicit_args(in_type, explicit_args)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sengthai.heng/Developer/test/.venv/lib/python3.12/site-packages/jax/_src/pjit.py", line 657, in _extract_implicit_args
    for d1, d2 in zip(aval.shape, arg.aval.shape):
                                  ^^^^^^^^
                                  ^^^^^^^^
AttributeError: 'tensor' object has no attribute 'aval'. Did you mean: 'ravel'?

System information

Name: PennyLane
Version: 0.41.1
Platform info:           macOS-15.5-arm64-arm-64bit
Python version:          3.12.0
Numpy version:           2.0.2
Scipy version:           1.15.1

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.