PyTorch-backed forward simulation

[rileyjmurray]

This PR introduces TorchForwardSimulator, a forward simulator (for computing circuit outcome probabilities) based on PyTorch. It uses automatic differentiation to compute the Jacobian of the map from model parameters to circuit outcome probabilities. In the future we could extend it to do computations on a system's GPU, or to use PyTorch-based optimization algorithms instead of pyGSTi's custom algorithms for MLE.

Approach

My approach required creating a new ModelMember subclass called Torchable. This subclass adds two required functions, called stateless_data and torch_base. Their meanings are given below:

pyGSTi/pygsti/modelmembers/torchable.py

Lines 18 to 43 in 1ec6909

	def stateless_data(self) -> Tuple:
	"""
	Return this ModelMember's data that is considered constant for purposes of model fitting.
	Note: the word "stateless" here is used in the sense of object-oriented programming.
	"""
	raise NotImplementedError()
	@staticmethod
	def torch_base(sd : Tuple, t_param : Tensor) -> Tensor:
	"""
	Suppose "obj" is an instance of some Torchable subclass. If we compute
	vec = obj.to_vector()
	t_param = torch.from_numpy(vec)
	sd = obj.stateless_data()
	t = type(obj).torch_base(sd, t_param)
	then t will be a PyTorch Tensor that represents "obj" in a canonical numerical way.
	The meaning of "canonical" is implementation dependent. If type(obj) implements
	the ``.base`` attribute, then a reasonable implementation will probably satisfy
	np.allclose(obj.base, t.numpy()).
	"""
	raise NotImplementedError()

In principle, TorchForwardSimulator can handle all models for which constituent parameterized ModelMembers are Torchable. So far I've only extended TPState, FullTPOp, and TPPOVM to be Torchable; these are the classes used in "full TP" GST.

The Python file that contains TorchForwardSimulator also defines two helper classes: StatelessCircuit and StatelessModel. I think it's fine to keep these classes as purely internal implementation-specific constructs for now. Depending on future performance optimizations of TorchForwardSimulator we might want to put them elsewhere in pyGSTi.

What should come after this PR

We should compare performance of TorchForwardSimulator to MapForwardSimulator on problems of interest. There's a chance that the former isn't faster than the latter with the current implementation. If that's the case then I should look at possible performance optimizations specifically inside TorchForwardSimulator.

We should add implementations of stateless_data and torch_base to GST models beyond "Full TP" (in particular I'd like to try CPTP).

Incidental changes

My implementation originally interacted with the following evotype classes

    <class 'pygsti.evotypes.densitymx[_slow].statereps.StateRepDense'>
    <class 'pygsti.evotypes.densitymx[_slow].opreps.OpRepDenseSuperop'>
    <class 'pygsti.evotypes.densitymx[_slow].effectreps.EffectRepConjugatedState'>

When I write [_slow] in the class names above you can put the empty string or just _slow, depending on the default evotype specified in evotypes.py.

To my surprise, I found that interacting with evotypes was neither necessary nor sufficient for what I wanted to accomplish. So while I did make changes in evotypes/densitymx_slow/ to remove unnecessary class inheritances and to add documentation, those changes were only to make life a little easier for future pyGSTi contributors.

[rileyjmurray]

This change is to resolve a deprecation warning.

[rileyjmurray]

This change is just to improve readability.

[sserita]

I'm giving a few quick comments here because I think this PR will actually take me some time to get through.

The main thing I am concerned and thinking about is the choice to extend the ModelMember class as opposed to adjusting the evotypes. The general purpose of the split between modelmembers and evotypes is so that we don't have to go through and implement these abstract methods in all modelmembers - we can make the change in the evotype and then any modelmember works.
It is totally possible that extending the API is the best way to do this. The TermSimulator follows a similar pattern where the API extensions are the cleanest way to do it. But I'll probably spend some time thinking about whether this is way we want to implement this. Probably a point of discussion for us in our dev meetings in the upcoming weeks.

[rileyjmurray]

@sserita, regarding

The main thing I am concerned and thinking about is the choice to extend the ModelMember class as opposed to adjusting the evotypes. The general purpose of the split between modelmembers and evotypes is so that we don't have to go through and implement these abstract methods in all modelmembers - we can make the change in the evotype and then any modelmember works.

Unfortunately there's no way to do this just through evotypes. Using pytorch's AD capabilities requires knowing the free parameters in a modelmember and how those free parameters map to the common parameterization-agnostic representation (i.e., representations in evotypes).

[rileyjmurray]

Try to use from __future__ import annotations to avoid use of TypeVars.

[rileyjmurray]

Notes from today's meeting:

• I removed empty base classes for OpRep, EffectRep, StateRep. I should make sure we aren't doing type checking elsewhere in pyGSTi that would be broken by this. (There's a chance we are and it simply isn't showing up in current tests.)

[rileyjmurray]

@coreyostrove, @sserita, @enielse: this is ready for review.

[rileyjmurray]

@coreyostrove, @sserita: is the dim ** 0.25 scale appropriate when the underlying vector space is something other than "tensor product of qubit space"?

[rileyjmurray]

@coreyostrove, @sserita: same question as above. Is the scale factor of dim ** -0.25 appropriate for general spaces?