class_docstr: Test the new VQE syntax of passing the Hamiltonian as an observable.
class_docstr_tok: ["Test","the","new","VQE","syntax","of","passing","the","Hamiltonian","as","an","observable","."]
docstr: Tests that distributed shots work with the new VQE design.
docstr_tok: ["Tests","that","distributed","shots","work","with","the","new","VQE","design","."]
code: def test_shot_distribution(self, shots, dim):
      """Tests that distributed shots work with the new VQE design."""
      import jax
  
      jax.config.update("jax_enable_x64", True)
  
      dev = qml.device("default.qubit", wires=2, shots=shots)
  
      @qml.qnode(dev)
      def circuit(weights, coeffs):
          qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1])
          H = qml.Hamiltonian(coeffs, obs)
          return qml.expval(H)
  
      obs = [qml.PauliZ(0), qml.PauliX(0) @ qml.PauliZ(1)]
      coeffs = np.array([0.1, 0.2])
      key = jax.random.PRNGKey(42)
      weights = jax.random.uniform(key, [2, 2, 3])
  
      res = circuit(weights, coeffs)
      grad = jax.jacobian(circuit, argnums=[1])(weights, coeffs)
      assert len(res) == dim
      assert qml.math.shape(grad) == (dim, 1, 2)
  
imports: import copy
  import numpy
  import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 3 for calling the TestNewVQE obj's underlying member method code with example usage: obj.test_shot_distribution(shots, dim) and returns: qml