code: def test_simple_circuit_with_prep():
      """Test that batching works for a simple circuit with a state preparation"""
      dev = qml.device("default.qubit", wires=2)
  
      @partial(qml.batch_input, argnum=1)
      @qml.qnode(dev)
      def circuit(inputs, weights):
          qml.StatePrep(np.array([0, 0, 1, 0]), wires=[0, 1])
          qml.RX(inputs, wires=0)
          qml.RY(weights[0], wires=0)
          qml.RY(weights[1], wires=1)
          return qml.expval(qml.PauliZ(1))
  
      batch_size = 5
      inputs = np.random.uniform(
          0, np.pi, (batch_size,), requires_grad=False
      )
      weights = np.random.uniform(-np.pi, np.pi, (2,))
  
      res = circuit(inputs, weights)
      assert res.shape == (batch_size,)
  
imports: from functools import partial
  import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: test_simple_circuit_with_prep() and returns: qml