[Return-types #11] Jax interface support shot vectors

[rmoyard]

Description of the Change:

This PR integrates shot vectors for QNode with JAX and derivates jacobian, hessian.

Benefits:

Users can use shot vectors everywhere with Jax QNode.

Example:

        import jax
        from jax.config import config

        config.update("jax_enable_x64", True)

        if diff_method == "adjoint":
            pytest.skip("Test does not supports adjoint because second order diff.")

        dev = qml.device(dev_name, wires=2, shots=(1, 10, 100))

        params = jax.numpy.array([0.1, 0.2])

        @qnode(dev, interface="jax", diff_method="parameter-shift", max_diff=2)
        def circuit(x):
            qml.RX(x[0], wires=[0])
            qml.RY(x[1], wires=[1])
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])

        hess = jax.hessian(circuit)(params)
        print(hess)

((DeviceArray([[ 0. , -0.5],
             [-0.5,  0. ]], dtype=float64), DeviceArray([[[-5.00000000e-01,  0.00000000e+00],
              [ 0.00000000e+00,  0.00000000e+00]],

             [[ 0.00000000e+00,  2.50000000e-01],
              [ 2.50000000e-01,  0.00000000e+00]],

             [[ 0.00000000e+00, -1.23259516e-32],
              [ 0.00000000e+00,  0.00000000e+00]],

             [[ 5.00000000e-01, -2.50000000e-01],
              [-2.50000000e-01,  0.00000000e+00]]], dtype=float64)), (DeviceArray([[-0.15 , -0.075],
             [-0.075,  0.   ]], dtype=float64), DeviceArray([[[-2.00000000e-01, -3.75000000e-02],
              [-3.75000000e-02,  0.00000000e+00]],

             [[-3.00000000e-01, -1.25000000e-02],
              [-1.25000000e-02,  0.00000000e+00]],

             [[ 3.75000000e-01,  5.00000000e-02],
              [ 5.00000000e-02,  0.00000000e+00]],

             [[ 1.25000000e-01, -5.55111512e-18],
              [ 0.00000000e+00,  0.00000000e+00]]], dtype=float64)), (DeviceArray([[-0.27 , -0.135],
             [-0.135, -0.26 ]], dtype=float64), DeviceArray([[[-0.325 , -0.0375],
              [-0.0375, -0.13  ]],

             [[-0.175 ,  0.035 ],
              [ 0.035 ,  0.13  ]],

             [[ 0.31  ,  0.0325],
              [ 0.0325,  0.    ]],

             [[ 0.19  , -0.03  ],
              [-0.03  ,  0.    ]]], dtype=float64)))

[antalszava]

@rmoyard this is amazing work! 🔥 Requesting changes because I think _shots_copies needs some improvements as commented, but the changes shouldn't be major. The PR otherwise is solid.

In addition, ideally, I think it would be best to separate the fix/improvement for compatibility with shot vector defs that contain tuples because it's more general than the JAX interface. Having said that, I see that making it stand-alone would require further testing.

I would be happy to help with this separation by adding unit tests, let me know. I think it's important that we separate such changes because otherwise they will be lost in some bigger scope PRs and it will be challenging to track them down.

Thanks for making the shot tuple fix btw!

[antalszava]

Hi @rmoyard, looks good! 🎉 Amazing work. 💯

Approving on the condition of:

• Shot vector tests being separated into a file;
• Mixing up the shots def to have the tuple be at the 1. pos (see suggestion).

The rest of the comments are more minor.

[eddddddy]

Hi @rmoyard, this is looking solid 💯

Just have a few minor questions, but otherwise this looks good on my end.