class_docstr: Tests for the SPSA gradient transform
class_docstr_tok: ["Tests","for","the","SPSA","gradient","transform"]
docstr: Make sure that custom samplers can be created as defined in the docs of spsa_grad.
docstr_tok: ["Make","sure","that","custom","samplers","can","be","created","as","defined","in","the","docs","of","spsa_grad","."]
code: def test_sampler_argument(self):
      """Make sure that custom samplers can be created as defined in the docs of spsa_grad."""
  
      def sampler_required_kwarg(
          indices, num_params, *args, rng
      ):  # pylint:disable=unused-argument
          direction = np.zeros(num_params)
          direction[indices] = rng.choice([-1, 0, 1], size=len(indices))
          return direction
  
      def sampler_required_arg_or_kwarg(
          indices, num_params, idx_rep, rng
      ):  # pylint:disable=unused-argument
          direction = np.zeros(num_params)
          direction[indices] = rng.choice([-1, 0, 1], size=len(indices))
          return direction
  
      def sampler_required_arg(
          indices, num_params, foo, idx_rep, rng, /
      ):  # pylint:disable=unused-argument
          """This should fail since spsa_grad passes rng as a kwarg."""
          direction = np.zeros(num_params)
          direction[indices] = rng.choice([-1, 0, 1], size=len(indices))
          return direction
  
      dev = qml.device("default.qubit", wires=1)
  
      tape = qml.tape.QuantumTape(
          [qml.RX(0.5, wires=0)], [qml.expval(qml.PauliZ(0))]
      )
  
      results = []
      for sampler in [
          sampler_required_arg_or_kwarg,
          sampler_required_kwarg,
      ]:
          sampler_rng = np.random.default_rng(42)
          tapes, proc_fn = spsa_grad(
              tape,
              sampler=sampler,
              num_directions=100,
              sampler_rng=sampler_rng,
          )
  
          res = qml.execute(tapes, dev)
          results.append(proc_fn(res))
  
      assert np.isclose(results[0], results[1], atol=0.1)
  
      err = "got some positional-only arguments passed as keyword arguments: 'rng'"
      with pytest.raises(TypeError, match=err):
          tapes, proc_fn = spsa_grad(
              tape,
              sampler=sampler_required_arg,
              num_directions=100,
              sampler_rng=sampler_rng,
          )
  
imports: import numpy
  import pytest
  import pennylane
  from pennylane import numpy
  from pennylane.devices import DefaultQubitLegacy
  from pennylane.gradients import spsa_grad
  from pennylane.gradients.spsa_gradient import _rademacher_sampler
  from pennylane.operation import AnyWires, Observable
label_desc: Use image node_id 1 for calling the TestSpsaGradient obj's underlying member method code with example usage: obj.test_sampler_argument() and returns: direction, direction, direction