Set default diff_method to adjoint for lightning devices (#1961)

**Context:**
Currently parameter-shift is used by default for differentiation, when
the `diff_method` is not provided. This is the most universal method,
but not the most performant, especially on lightning devices which
support adjoint differentiation (for expval)

**Description of the Change:**
For lightning devices, set `diff_method` to `adjoint` when it is set to
`best`, which is the default argument.

Note: currently only a single expval is supported - this needs to be
fixed. See
[discussion](#1961 (comment))

**Benefits:**
Significantly faster differentiation performance on lightning devices.

**Possible Drawbacks:**

**Related GitHub Issues:**
#1171

[sc-74760]

Author: josephleekl

doc/releases/changelog-dev.md

@@ -4,6 +4,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* Adjoint differentiation is used by default when executing on lightning devices, significantly reduces gradient computation time.
+  [(#1961)](https://github.com/PennyLaneAI/catalyst/pull/1961)
+
 * Added `detensorizefunctionboundary` pass to remove scalar tensors across function boundaries and enabled `symbol-dce` pass to remove dead functions, reducing the number of instructions for compilation.
   [(#1904)](https://github.com/PennyLaneAI/catalyst/pull/1904)
 

frontend/catalyst/jax_primitives_utils.py

@@ -23,11 +23,7 @@
 from jax.interpreters import mlir
 from jaxlib.mlir.dialects.builtin import ModuleOp
 from jaxlib.mlir.dialects.func import CallOp
-from mlir_quantum.dialects._transform_ops_gen import (
-    ApplyRegisteredPassOp,
-    NamedSequenceOp,
-    YieldOp,
-)
+from mlir_quantum.dialects._transform_ops_gen import ApplyRegisteredPassOp, NamedSequenceOp, YieldOp
 from mlir_quantum.dialects.catalyst import LaunchKernelOp
 
 
@@ -113,13 +109,30 @@ def lower_callable_to_funcop(ctx, callable_, call_jaxpr):
 
     if isinstance(callable_, qml.QNode):
         func_op.attributes["qnode"] = ir.UnitAttr.get()
-        # "best", the default option in PennyLane, chooses backprop on the device
-        # if supported and parameter-shift otherwise. Emulating the same behaviour
-        # would require generating code to query the device.
-        # For simplicity, Catalyst instead defaults to parameter-shift.
-        diff_method = (
-            "parameter-shift" if callable_.diff_method == "best" else str(callable_.diff_method)
-        )
+
+        diff_method = str(callable_.diff_method)
+
+        if diff_method == "best":
+
+            def only_single_expval():
+                found_expval = False
+                for eqn in call_jaxpr.eqns:
+                    name = eqn.primitive.name
+                    if name in {"probs", "counts", "sample"}:
+                        return False
+                    elif name == "expval":
+                        if found_expval:
+                            return False
+                        found_expval = True
+                return True
+
+            device_name = getattr(getattr(callable_, "device", None), "name", None)
+
+            if device_name and "lightning" in device_name and only_single_expval():
+                diff_method = "adjoint"
+            else:
+                diff_method = "parameter-shift"
+
         func_op.attributes["diff_method"] = ir.StringAttr.get(diff_method)
 
     return func_op