class_docstr: Autodiff tests for all frameworks
class_docstr_tok: ["Autodiff","tests","for","all","frameworks"]
docstr: Test that results after splitting are still differentiable with jax
  with multiple parameters
docstr_tok: ["Test","that","results","after","splitting","are","still","differentiable","with","jax","with","multiple","parameters"]
code: def test_split_with_jax_multi_params_probs(self):
      """Test that results after splitting are still differentiable with jax
      with multiple parameters"""
  
      import jax
      import jax.numpy as jnp
  
      dev = qml.device("default.qubit.jax", wires=2)
  
      @qml.qnode(dev)
      def circuit(x, y):
          qml.RX(x, wires=0)
          qml.RY(y, wires=1)
          return (
              qml.probs(wires=[0, 1]),
              qml.expval(qml.PauliX(0) @ qml.PauliX(1)),
          )
  
      x = jnp.array(0.5)
      y = jnp.array(0.5)
  
      res = split_non_commuting(circuit)(x, y)
      res = jnp.concatenate([res[0]] + [jnp.stack(res[1:])], axis=0)
      assert all(np.isclose(res, exp_res_probs))
  
      grad = jax.jacobian(split_non_commuting(circuit), argnums=[0, 1])(
          x, y
      )
  
      assert isinstance(grad, tuple)
      assert len(grad) == 2
  
      for meas_grad in grad:
          assert isinstance(meas_grad, tuple)
          assert len(meas_grad) == 2
          assert all(isinstance(g, jnp.ndarray) for g in meas_grad)
  
      # reshape the returned gradient to the right shape
      grad = jnp.concatenate(
          [
              jnp.concatenate(
                  [
                      grad[0][0].reshape(-1, 1),
                      grad[0][1].reshape(-1, 1),
                  ],
                  axis=1,
              ),
              jnp.concatenate(
                  [
                      grad[1][0].reshape(-1, 1),
                      grad[1][1].reshape(-1, 1),
                  ],
                  axis=1,
              ),
          ],
          axis=0,
      )
      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 6 for calling the TestAutodiffSplitNonCommuting obj's underlying member method code with example usage: obj.test_split_with_jax_multi_params_probs() and returns: qml,  qml