Repository: PennyLaneAI/pennylane

Issue:
[BUG] Possible bug in Hamiltonian observable unflattening when using `jax.jit`
### Expected behavior

```python
import pennylane as qml
from pennylane import numpy as pnp
import numpy as np
import pytest
import jax

class TestBroadcastExpand:
    def test_expansion_qnode(self):

        coeffs0 = [0.3, -5.1]
        obs = qml.Hamiltonian(qml.math.array(coeffs0), [qml.PauliZ(0), qml.PauliY(1)])
        @qml.qnode(qml.device("default.qubit", wires=2), interface="jax")
        def circuit(obs):
            return qml.state()

        expected = circuit(obs)
        print(expected)
```

Should be compatible with `jax.jit`

### Actual behavior

Unflattening observable Hamiltonians when `jax.jit`ing a function raises an exception.

```python
import pennylane as qml
from pennylane import numpy as pnp
import numpy as np
import pytest
import jax

class TestBroadcastExpand:
    def test_expansion_qnode(self):

        coeffs0 = [0.3, -5.1]
        obs = qml.Hamiltonian(qml.math.array(coeffs0), [qml.PauliZ(0), qml.PauliY(1)])
        @qml.qnode(qml.device("default.qubit", wires=2), interface="jax")
        def circuit(obs):
            return qml.state()

        expected = circuit(obs)
        print(expected)
        observed = jax.jit(circuit)(obs) # <-- extra line causes an exception
```

### Additional information

I run this file with pytest.

### Source code

```python
import pennylane as qml
from pennylane import numpy as pnp
import numpy as np
import pytest
import jax

class TestBroadcastExpand:
    def test_expansion_qnode(self):

        coeffs0 = [0.3, -5.1]
        obs = qml.Hamiltonian(qml.math.array(coeffs0), [qml.PauliZ(0), qml.PauliY(1)])
        @qml.qnode(qml.device("default.qubit", wires=2), interface="jax")
        def circuit(obs):
            return qml.state()

        expected = circuit(obs) # <-- succeeds
        print(expected)
        observed = jax.jit(circuit)(obs) # <-- raises exception
```


### Tracebacks

```shell
env/lib/python3.10/site-packages/jax/_src/pjit.py:253: in cache_miss
    outs, out_flat, out_tree, args_flat, jaxpr = _python_pjit_helper(
env/lib/python3.10/site-packages/jax/_src/pjit.py:161: in _python_pjit_helper
    args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
env/lib/python3.10/site-packages/jax/_src/api.py:324: in infer_params
    return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
env/lib/python3.10/site-packages/jax/_src/pjit.py:491: in common_infer_params
    jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
env/lib/python3.10/site-packages/jax/_src/pjit.py:969: in _pjit_jaxpr
    jaxpr, final_consts, out_type = _create_pjit_jaxpr(
env/lib/python3.10/site-packages/jax/_src/linear_util.py:345: in memoized_fun
    ans = call(fun, *args)
env/lib/python3.10/site-packages/jax/_src/pjit.py:922: in _create_pjit_jaxpr
    jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
env/lib/python3.10/site-packages/jax/_src/profiler.py:314: in wrapper
    return func(*args, **kwargs)
env/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py:2155: in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
env/lib/python3.10/site-packages/jax/_src/interpreters/partial_eval.py:2177: in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
env/lib/python3.10/site-packages/jax/_src/linear_util.py:183: in call_wrapped
    args, kwargs = next(gen)
env/lib/python3.10/site-packages/jax/_src/api_util.py:86: in flatten_fun_nokwargs
    py_args = tree_unflatten(in_tree, args_flat)
env/lib/python3.10/site-packages/jax/_src/tree_util.py:93: in tree_unflatten
    return treedef.unflatten(leaves)
../pennylane/pennylane/pytrees.py:35: in jax_unflatten
    return unflatten_fn(parameters, aux)
../pennylane/pennylane/ops/qubit/hamiltonian.py:178: in _unflatten
    new_op = cls(data[0], data[1])
../pennylane/pennylane/ops/qubit/hamiltonian.py:191: in __init__
    if qml.math.shape(coeffs)[0] != len(observables):
env/lib/python3.10/site-packages/autoray/autoray.py:79: in do
    return get_lib_fn(backend, fn)(*args, **kwargs)
../pennylane/pennylane/math/single_dispatch.py:39: in <lambda>
    ar.register_function("builtins", "shape", lambda x: np.array(x).shape)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = Traced<ShapedArray(float32[])>with<DynamicJaxprTrace(level=1/0)>, args = (), kw = {}

    def __array__(self, *args, **kw):
>     raise TracerArrayConversionError(self)
E     jax._src.traceback_util.UnfilteredStackTrace: jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float32[].
E     The error occurred while tracing the function circuit at /home/erick.ochoalopez/Code/catalyst-latest/bug.py:12 for jit. This concrete value was not available in Python because it depends on the value of the argument obs[<flat index 0>][0].
E     See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
E     
E     The stack trace below excludes JAX-internal frames.
E     The preceding is the original exception that occurred, unmodified.
E     
E     --------------------

env/lib/python3.10/site-packages/jax/_src/core.py:611: UnfilteredStackTrace

The above exception was the direct cause of the following exception:

self = <bug.TestBroadcastExpand object at 0x7f6f3e85b4f0>

    def test_expansion_qnode(self):
    
        coeffs0 = [0.3, -5.1]
        obs = qml.Hamiltonian(qml.math.array(coeffs0), [qml.PauliZ(0), qml.PauliY(1)])
        @qml.qnode(qml.device("default.qubit", wires=2), interface="jax")
        def circuit(obs):
            return qml.state()
    
        expected = circuit(obs)
        print(expected)
>       observed = jax.jit(circuit)(obs)

bug.py:18: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../pennylane/pennylane/pytrees.py:35: in jax_unflatten
    return unflatten_fn(parameters, aux)
../pennylane/pennylane/ops/qubit/hamiltonian.py:178: in _unflatten
    new_op = cls(data[0], data[1])
../pennylane/pennylane/ops/qubit/hamiltonian.py:191: in __init__
    if qml.math.shape(coeffs)[0] != len(observables):
env/lib/python3.10/site-packages/autoray/autoray.py:79: in do
    return get_lib_fn(backend, fn)(*args, **kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

x = (Traced<ShapedArray(float32[])>with<DynamicJaxprTrace(level=1/0)>, Traced<ShapedArray(float32[])>with<DynamicJaxprTrace(level=1/0)>)

>   ar.register_function("builtins", "shape", lambda x: np.array(x).shape)
E   jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float32[].
E   The error occurred while tracing the function circuit at /home/erick.ochoalopez/Code/catalyst-latest/bug.py:12 for jit. This concrete value was not available in Python because it depends on the value of the argument obs[<flat index 0>][0].
E   See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

../pennylane/pennylane/math/single_dispatch.py:39: TracerArrayConversionError
```


### System information

```shell
Name: PennyLane
Version: 0.33.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/erick.ochoalopez/Code/catalyst-latest/env/lib/python3.10/site-packages
Editable project location: /home/erick.ochoalopez/Code/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: pennylane-catalyst, PennyLane-Lightning

Platform info:           Linux-5.15.0-83-generic-x86_64-with-glibc2.31
Python version:          3.10.11
Numpy version:           1.23.5
Scipy version:           1.10.0
Installed devices:
- default.gaussian (PennyLane-0.33.0.dev0)
- default.mixed (PennyLane-0.33.0.dev0)
- default.qubit (PennyLane-0.33.0.dev0)
- default.qubit.autograd (PennyLane-0.33.0.dev0)
- default.qubit.jax (PennyLane-0.33.0.dev0)
- default.qubit.legacy (PennyLane-0.33.0.dev0)
- default.qubit.tf (PennyLane-0.33.0.dev0)
- default.qubit.torch (PennyLane-0.33.0.dev0)
- default.qutrit (PennyLane-0.33.0.dev0)
- null.qubit (PennyLane-0.33.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.32.0)


# Using this commit-id for my version of pennylane, but this issue is likely unrelated to the commit below.

cbc8d19e6a36f0d4d38b44eabd4a8b98db677327
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
--- a/pennylane/math/__init__.py
+++ b/pennylane/math/__init__.py
@@ -85,6 +85,7 @@
     cast,
     cast_like,
     convert_like,
+    get_deep_interface,
     get_interface,
     in_backprop,
     is_abstract,
@@ -152,6 +153,7 @@ def __getattr__(name):
     "frobenius_inner_product",
     "get_dtype_name",
     "get_interface",
+    "get_deep_interface",
     "get_trainable_indices",
     "in_backprop",
     "is_abstract",
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -19,6 +19,9 @@
 import autoray as ar
 import numpy as np
 import semantic_version
+from scipy.linalg import block_diag as _scipy_block_diag
+
+from .utils import get_deep_interface
 
 
 def _i(name):
@@ -35,8 +38,22 @@ def _i(name):
 
 # ------------------------------- Builtins -------------------------------- #
 
-ar.register_function("builtins", "ndim", lambda x: np.ndim(np.array(x)))
-ar.register_function("builtins", "shape", lambda x: np.array(x).shape)
+
+def _builtins_ndim(x):
+    interface = get_deep_interface(x)
+    x = ar.numpy.asarray(x, like=interface)
+    return ar.ndim(x)
+
+
+def _builtins_shape(x):
+    interface = get_deep_interface(x)
+    x = ar.numpy.asarray(x, like=interface)
+    return ar.shape(x)
+
+
+ar.register_function("builtins", "ndim", _builtins_ndim)
+ar.register_function("builtins", "shape", _builtins_shape)
+
 
 # -------------------------------- SciPy --------------------------------- #
 # the following is required to ensure that SciPy sparse Hamiltonians passed to
@@ -50,7 +67,6 @@ def _i(name):
 
 
 # -------------------------------- NumPy --------------------------------- #
-from scipy.linalg import block_diag as _scipy_block_diag
 
 ar.register_function("numpy", "flatten", lambda x: x.flatten())
 ar.register_function("numpy", "coerce", lambda x: x)
diff --git a/pennylane/math/utils.py b/pennylane/math/utils.py
--- a/pennylane/math/utils.py
+++ b/pennylane/math/utils.py
@@ -279,6 +279,39 @@ def _get_interface_of_single_tensor(tensor):
     return res
 
 
+def get_deep_interface(value):
+    """
+    Given a deep data structure with interface-specific scalars at the bottom, return their
+    interface name.
+
+    Args:
+        value (list, tuple): A deep list-of-lists, tuple-of-tuples, or combination with
+            interface-specific data hidden within it
+
+    Returns:
+        str: The name of the interface deep within the value
+
+    **Example**
+
+    >>> x = [[jax.numpy.array(1), jax.numpy.array(2)], [jax.numpy.array(3), jax.numpy.array(4)]]
+    >>> get_deep_interface(x)
+    'jax'
+
+    This can be especially useful when converting to the appropriate interface:
+
+    >>> qml.math.asarray(x, like=qml.math.get_deep_interface(x))
+    Array([[1, 2],
+       [3, 4]], dtype=int64)
+
+    """
+    itr = value
+    while isinstance(itr, (list, tuple)):
+        if len(itr) == 0:
+            return "builtins"
+        itr = itr[0]
+    return ar.infer_backend(itr)
+
+
 def is_abstract(tensor, like=None):
     """Returns True if the tensor is considered abstract.
 


Test patch:
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -134,6 +134,12 @@ def test_return_numpy_box(self):
         res = fn.get_interface(y, x)
         assert res == "numpy"
 
+    def test_get_deep_interface(self):
+        """Test get_deep_interface returns the interface of deep values."""
+        assert fn.get_deep_interface([()]) == "builtins"
+        assert fn.get_deep_interface(([1, 2], [3, 4])) == "builtins"
+        assert fn.get_deep_interface([[jnp.array(1.1)]]) == "jax"
+
 
 test_abs_data = [
     (1, -2, 3 + 4j),
@@ -1327,6 +1333,13 @@ def test_shape(shape, interface, create_array):
     assert fn.shape(t) == shape
 
 
+@pytest.mark.parametrize("interface", ["numpy", "autograd", "jax", "torch", "tensorflow"])
+def test_shape_and_ndim_deep(interface):
+    val = [[fn.asarray(1, like=interface)]]
+    assert fn.shape(val) == (1, 1)
+    assert fn.ndim(val) == 2
+
+
 @pytest.mark.parametrize(
     "x, expected",
     (