code: def set_decomposition(custom_decomps, dev, decomp_depth=10):
      """Context manager for setting custom decompositions.
  
      Args:
          custom_decomps (Dict[Union(str, qml.operation.Operation), Callable]): Custom
              decompositions to be applied by the device at runtime.
          dev (pennylane.Device): A quantum device.
          decomp_depth: The maximum depth of the expansion.
  
      **Example**
  
      Suppose we would like a custom expansion function that decomposes all CNOTs
      into CZs. We first define a decomposition function:
  
      .. code-block:: python
  
          def custom_cnot(wires):
              return [
                  qml.Hadamard(wires=wires[1]),
                  qml.CZ(wires=[wires[0], wires[1]]),
                  qml.Hadamard(wires=wires[1])
              ]
  
      This context manager can be used to temporarily change a devices expansion
      function to one that takes into account the custom decompositions.
  
      .. code-block:: python
  
          dev = qml.device("default.qubit", wires=2)
  
          @qml.qnode(dev, expansion_strategy="device")
          def circuit():
              qml.CNOT(wires=[0, 1])
              return qml.expval(qml.PauliZ(wires=0))
  
      >>> print(qml.draw(circuit)())
      0: ─╭●─┤  <Z>
      1: ─╰X─┤
  
      Now let's set up a context where the custom decomposition will be applied:
  
      >>> with qml.transforms.set_decomposition({qml.CNOT : custom_cnot}, dev):
      ...     print(qml.draw(circuit, wire_order=[0, 1])())
      0: ────╭●────┤  <Z>
      1: ──H─╰Z──H─┤
  
      """
      if isinstance(dev, qml.Device):
          original_custom_expand_fn = dev.custom_expand_fn
  
          # Create a new expansion function; stop at things that do not have
          # custom decompositions, or that satisfy the regular device stopping criteria
          new_custom_expand_fn = create_decomp_expand_fn(
              custom_decomps, dev, decomp_depth=decomp_depth
          )
  
          # Set the custom expand function within this context only
          try:
              dev.custom_expand(new_custom_expand_fn)
              yield
  
          finally:
              dev.custom_expand_fn = original_custom_expand_fn
  
      else:
          original_preprocess = dev.preprocess
          new_preprocess = _create_decomp_preprocessing(
              custom_decomps, dev, decomp_depth=decomp_depth
          )
  
          try:
              dev.preprocess = new_preprocess
              yield
  
          finally:
              dev.preprocess = original_preprocess
  
file_path: pennylane/pennylane/transforms/tape_expand.py
imports: import contextlib
  import pennylane
  from pennylane.operation import has_gen, gen_is_multi_term_hamiltonian, has_grad_method, has_nopar, has_unitary_gen, is_measurement, is_trainable, not_tape
label_desc: Use image node_id 7 for calling a global function with example usage: set_decomposition(custom_decomps, dev, decomp_depth) without return types