[BUG] qml.ctrl is sometimes promoting dtypes when used with JAX

[CatalinaAlbornoz]

Expected behavior

qml.ctrl maintains the user-defined dtype

Actual behavior

qml.ctrl is sometimes promoting dtypes when used with JAX

Additional information

It looks like autoray turns all Python complex values to complex128 values when using autoray.to_numpy().

Originally reported here in the Forum.

If we remove one qml.ctrl from the code below we get no warning.

Source code

import jax
import jax.numpy as jnp
import pennylane as qml

jax.config.update("jax_enable_x64", False) 

@qml.qnode(qml.device('default.qubit', wires=range(3)))
def circuit(phi, theta, omega):
    qml.ctrl(qml.ctrl(qml.Rot, control=(1), control_values=1), control=(0), control_values=1)(phi, theta, omega, wires=2) 
    return qml.expval(qml.Z(0))

phi = theta = omega = jnp.array(1.2)
qml.draw_mpl(circuit)(phi, theta, omega); # draw the circuit
circuit(phi, theta, omega)

Tracebacks

/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py:118: UserWarning: Explicitly requested dtype <class 'jax.numpy.complex128'> requested in astype is not available, and will be truncated to dtype complex64. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.
  return lax_numpy.astype(self, dtype, copy=copy, device=device)
/usr/local/lib/python3.10/dist-packages/jax/_src/numpy/array_methods.py:118: UserWarning: Explicitly requested dtype <class 'jax.numpy.complex128'> requested in astype is not available, and will be truncated to dtype complex64. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.
  return lax_numpy.astype(self, dtype, copy=copy, device=device)
Array(1., dtype=float32)

System information

Name: PennyLane
Version: 0.39.0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /usr/local/lib/python3.10/dist-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-6.1.85+-x86_64-with-glibc2.35
Python version:          3.10.12
Numpy version:           1.26.4
Scipy version:           1.13.1
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.39.0)
- default.clifford (PennyLane-0.39.0)
- default.gaussian (PennyLane-0.39.0)
- default.mixed (PennyLane-0.39.0)
- default.qubit (PennyLane-0.39.0)
- default.qutrit (PennyLane-0.39.0)
- default.qutrit.mixed (PennyLane-0.39.0)
- default.tensor (PennyLane-0.39.0)
- null.qubit (PennyLane-0.39.0)
- reference.qubit (PennyLane-0.39.0)

Existing GitHub issues

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