class_docstr: Autodiff tests for all frameworks
class_docstr_tok: ["Autodiff","tests","for","all","frameworks"]
docstr: Test resulting after splitting non-commuting tapes with expval and probs measurements
  are still differentiable with jax
docstr_tok: ["Test","resulting","after","splitting","non-commuting","tapes","with","expval","and","probs","measurements","are","still","differentiable","with","jax"]
code: def test_split_with_jax_probs(self):
      """Test resulting after splitting non-commuting tapes with expval and probs measurements
      are still differentiable with jax"""
      import jax
      import jax.numpy as jnp
  
      dev = qml.device("default.qubit.jax", wires=2)
  
      @qml.qnode(dev)
      def circuit(params):
          qml.RX(params[0], wires=0)
          qml.RY(params[1], wires=1)
          return (
              qml.probs(wires=[0, 1]),
              qml.expval(qml.PauliX(0) @ qml.PauliX(1)),
          )
  
      params = jnp.array([0.5, 0.5])
      res = split_non_commuting(circuit)(params)
      res = jnp.concatenate([res[0]] + [jnp.stack(res[1:])], axis=0)
  
      grad = jax.jacobian(split_non_commuting(circuit))(params)
      grad = jnp.concatenate([grad[0]] + [jnp.stack(grad[1:])], axis=0)
  
      assert all(np.isclose(res, exp_res_probs, atol=0.05))
      assert all(np.isclose(grad, exp_grad_probs, atol=0.05).flatten())
  
imports: import itertools
  import pytest
  import numpy
  import pennylane
  import pennylane.numpy
  from pennylane.transforms import split_non_commuting
label_desc: Use image node_id 4 for calling the TestAutodiffSplitNonCommuting obj's underlying member method code with example usage: obj.test_split_with_jax_probs() and returns: qml,  qml