chore: clean-up type hints in qnode.py and resolution.py (#8086)

**Context:**

Fix some annoying `mypy` complaints. 

**Benefits:** Better code.

Author: andrijapau

doc/releases/changelog-dev.md

@@ -462,6 +462,9 @@
 
 <h3>Internal changes ⚙️</h3>
 
+* Improve type hinting internally.
+  [(#8086)](https://github.com/PennyLaneAI/pennylane/pull/8086)
+
 * The `cond` primitive with program capture no longer stores missing false branches as `None`, instead storing them
   as jaxprs with no output.
   [(#8080)](https://github.com/PennyLaneAI/pennylane/pull/8080)

pennylane/workflow/qnode.py

@@ -44,14 +44,16 @@
 logger.addHandler(logging.NullHandler())
 
 if TYPE_CHECKING:
+    from typing import TypeAlias
+
     from pennylane.concurrency.executors import ExecBackends
     from pennylane.devices import Device, LegacyDevice
     from pennylane.math import SupportedInterfaceUserInput
     from pennylane.transforms.core import TransformContainer
     from pennylane.typing import Result
     from pennylane.workflow.resolution import SupportedDiffMethods
 
-    SupportedDeviceAPIs = LegacyDevice | Device
+    SupportedDeviceAPIs: TypeAlias = LegacyDevice | Device
 
 
 def _convert_to_interface(result, interface: Interface):
@@ -119,7 +121,10 @@ def _to_qfunc_output_type(results: Result, qfunc_output, has_partitioned_shots)
     return pytrees.unflatten(results, qfunc_output_structure)
 
 
-def _validate_mcm_config(postselect_mode: str, mcm_method: str) -> None:
+def _validate_mcm_config(
+    postselect_mode: Literal["hw-like", "fill-shots"] | None,
+    mcm_method: Literal["deferred", "one-shot", "tree-traversal"] | None,
+) -> None:
     qml.devices.MCMConfig(postselect_mode=postselect_mode, mcm_method=mcm_method)
 
 
@@ -581,7 +586,7 @@ def __init__(
         self.diff_method = diff_method
         _validate_diff_method(self.device, self.diff_method)
 
-        self.capture_cache = LRUCache(maxsize=1000)
+        self.capture_cache: LRUCache = LRUCache(maxsize=1000)
         if isinstance(static_argnums, int):
             static_argnums = (static_argnums,)
         self.static_argnums = sorted(static_argnums)
@@ -727,7 +732,7 @@ def circuit(x):
         original_init_args["gradient_kwargs"] = original_init_args["gradient_kwargs"] or {}
         # nested dictionary update
         new_gradient_kwargs = kwargs.pop("gradient_kwargs", {})
-        old_gradient_kwargs = original_init_args.get("gradient_kwargs").copy()
+        old_gradient_kwargs = (original_init_args.get("gradient_kwargs", {})).copy()
         old_gradient_kwargs.update(new_gradient_kwargs)
         kwargs["gradient_kwargs"] = old_gradient_kwargs
 

pennylane/workflow/resolution.py

@@ -15,12 +15,11 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
 from copy import copy
 from dataclasses import replace
 from importlib.metadata import version
 from importlib.util import find_spec
-from typing import TYPE_CHECKING, Literal, get_args
+from typing import TYPE_CHECKING, Any, Literal, get_args
 from warnings import warn
 
 from packaging.version import Version
@@ -146,7 +145,7 @@ def _resolve_mcm_config(
 ) -> qml.devices.MCMConfig:
     """Helper function to resolve the mid-circuit measurements configuration based on
     execution parameters"""
-    updated_values = {}
+    updated_values: dict[str, Any] = {}
 
     if not finite_shots:
         updated_values["postselect_mode"] = None
@@ -287,10 +286,10 @@ def _resolve_execution_config(
     Returns:
         qml.devices.ExecutionConfig: resolved execution configuration
     """
-    updated_values = {}
+    updated_values: dict[str, Any] = {}
 
-    if execution_config.interface in {Interface.JAX, Interface.JAX_JIT} and not isinstance(
-        execution_config.gradient_method, Callable
+    if execution_config.interface in {Interface.JAX, Interface.JAX_JIT} and not callable(
+        execution_config.gradient_method
     ):
         updated_values["grad_on_execution"] = False
     execution_config = _resolve_diff_method(execution_config, device, tape=tapes[0])