class_docstr: Test that differentiating tapes with ``pulse_odegen`` works.
class_docstr_tok: ["Test","that","differentiating","tapes","with","``","pulse_odegen","``","works","."]
docstr: Test that a single pulse with multiple Hamiltonian terms is
  differentiated correctly when setting ``argnum``.
docstr_tok: ["Test","that","a","single","pulse","with","multiple","Hamiltonian","terms","is","differentiated","correctly","when","setting","``","argnum","``","."]
code: def test_single_pulse_multi_term_argnum(self, dev_name, argnum):
      """Test that a single pulse with multiple Hamiltonian terms is
      differentiated correctly when setting ``argnum``."""
      import jax
      import jax.numpy as jnp
  
      dev = qml.device(dev_name, wires=1)
  
      H = jnp.polyval * X(0) + qml.pulse.constant * X(0)
      x = jnp.array([0.4, 0.2, 0.1])
      y = jnp.array(0.6)
      t = [0.2, 0.3]
      op = qml.evolve(H)([x, y], t=t)
      tape = qml.tape.QuantumScript([op], [qml.expval(Z(0))])
  
      val = qml.execute([tape], dev)
      theta = integral_of_polyval(x, t) + y * (t[1] - t[0])
      assert qml.math.allclose(val, jnp.cos(2 * theta))
  
      # Argnum=[0] or 0
      par_jac_0 = jax.jacobian(integral_of_polyval)(x, t)
      exp_grad_0 = -2 * par_jac_0 * jnp.sin(2 * theta)
      # Argnum=[1] or 1
      par_jac_1 = t[1] - t[0]
      exp_grad_1 = -2 * par_jac_1 * jnp.sin(2 * theta)
      exp_grads = [exp_grad_0, exp_grad_1]
  
      _tapes, fn = pulse_odegen(tape, argnum=argnum)
      assert len(_tapes) == 2
  
      grad = fn(qml.execute(_tapes, dev))
      assert isinstance(grad, tuple) and len(grad) == 2
      assert (
          isinstance(grad[0], jnp.ndarray) and grad[0].shape == x.shape
      )
      assert (
          isinstance(grad[1], jnp.ndarray) and grad[1].shape == y.shape
      )
      argnum_int = argnum if isinstance(argnum, int) else argnum[0]
      assert qml.math.allclose(grad[argnum_int], exp_grads[argnum_int])
      assert qml.math.allclose(grad[1 - argnum_int], 0.0)
  
imports: import copy
  import pytest
  import numpy
  import pennylane
  from pennylane import numpy
  from pennylane.ops.qubit.special_unitary import pauli_basis_matrices, pauli_basis_strings
  from pennylane.math import expand_matrix
  from pennylane.gradients.pulse_gradient_odegen import _generate_tapes_and_coeffs, pulse_odegen, _insert_op, _nonzero_coeffs_and_words, _one_parameter_generators, _one_parameter_paulirot_coeffs, _parshift_and_contract
label_desc: Use image node_id 3 for calling the TestPulseOdegenTape obj's underlying member method code with example usage: obj.test_single_pulse_multi_term_argnum(dev_name, argnum) without return types