Fix and test Jax automatic differentiation #14
Conversation
src/qiboml/operations/expectation.py
Outdated
| return _with_tf(**kwargs) | ||
|
|
||
| if isinstance(qibo_backend, JaxBackend): | ||
| return _with_jax(**kwargs) |
There was a problem hiding this comment.
maybe the _with_x functions should directly be part of the backend, this way here you don't need to check which backend is used but just do qibo_backend._with(...)
There was a problem hiding this comment.
To be fair, it would be even nicer if I could construct a Qibo observable without any specific backend...
There was a problem hiding this comment.
We need a way to detect which is the frontend, and then build the differentiation rule on top of this. To have an agnostic observable doesn't seems to me a good idea, because it is the way we detect the frontend from the user code.
Regarding the _with(...) suggestion, it could be indeed a nice idea, but then I see a problem: this function will be a third version of the expectation we have in the backends. Moreover, it is a strange expectation since it wants as an argument a differentiation_rule, which (my opinion) makes sense only in this QML context.
src/qiboml/operations/expectation.py
Outdated
| for p in range(len(params)): | ||
| gradients.append( | ||
| differentiation_rule( | ||
| circuit=circuit, | ||
| hamiltonian=observable, | ||
| parameter_index=p, | ||
| initial_state=initial_state, | ||
| nshots=nshots, | ||
| backend=backend, | ||
| ) | ||
| ) |
There was a problem hiding this comment.
Batched execution wold be great here, for example, as you could perform psr on several different parameters parallely. Namely the batch would be
((theta_0 + eps, theta_1, ..., theta_n),
(theta_0, theta_1 + eps, ..., theta_n),
...,
(theta_0, theta_1, ..., theta_n + eps))
| def _expectation(params): | ||
| params = jax.numpy.array(params) | ||
|
|
||
| def grad(params): |
There was a problem hiding this comment.
How about jitting grad? Is it possible?
There was a problem hiding this comment.
It is, but in principle it could be done even later on: JAX transformations are composable, so you could do grad(grad(jit(grad(...)))).
(but I'm not sure which is the best way for performances, if it's better to jit each building block or rely on the external compilation - also, in this case, I'm not sure if there is any benefit in jitting, since it could be the result of hardware execution...)
There was a problem hiding this comment.
I jitted all the functions we are returning now, but I think we should build a Jax version of the Qibojit backend, with a more fine-grained jitting, to boost the performances.
| @@ -14,7 +16,7 @@ def parameter_shift( | |||
| initial_state: Optional[Union[np.ndarray, qibo.Circuit]] = None, | |||
| scale_factor: float = 1.0, | |||
| nshots: int = None, | |||
| backend: str = "qibojit", | |||
| exec_backend: qibo.backends.Backend = NumbaBackend, | |||
There was a problem hiding this comment.
| exec_backend: qibo.backends.Backend = NumbaBackend, | |
| exec_backend: qibo.backends.Backend = NumbaBackend(), |
I think this should be NumbaBackend() otherwise you'll be passing just the class. In general I am not against using numba by default, but then qibojit has to become a mandatory dependency of qiboml. Personally, I think the JaxBackend is the best candidate to be used as default, if not the only one at all.
| def Unitary(self, u): | ||
| return self._cast(u, dtype=self.dtype) | ||
|
|
||
|
|
||
| class JaxBackend(NumpyBackend): |
There was a problem hiding this comment.
I would probably rewrite the execute_circuit and apply_gate functions of the backend to properly make use of the jit decorator
No description provided.