code: def cost_fn_single(x):
      first_term = qml.math.convert_like(qml.math.diag([1.0, 0]), x)
      second_term = qml.math.convert_like(qml.math.diag([0, 1.0]), x)
  
      x = qml.math.cast_like(x, first_term)
      if len(qml.math.shape(x)) == 0:
          state1 = (
              qml.math.cos(x / 2) ** 2 * first_term
              + qml.math.sin(x / 2) ** 2 * second_term
          )
      else:
          # broadcasting
          x = x[:, None, None]
          state1 = (
              qml.math.cos(x / 2) ** 2 * first_term
              + qml.math.sin(x / 2) ** 2 * second_term
          )
  
      state2 = qml.math.convert_like(qml.math.diag([1, 0]), state1)
  
      return qml.math.fidelity(state1, state2) + qml.math.fidelity(
          state2, state1
      )
  
imports: import numpy
  import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: cost_fn_single(x) and returns: unknown