Change TransformState to NamedTuple

[SamDuffield]

This is quite a large PR unfortunately but fixes #83 and also cleans up the TransformState code by moving from mutable dataclass to immutable NamedTuple which is also better encapsulation practice for functional code.

One thing that's nice is that NamedTuple is already added to the pytree registry for optree and torch (fixing #83), so we don't need to do that manually as before.

There was an issue with the aux handling as modifying aux is not possible to do in-place as we don't know the structure of aux before log_posterior is called, also aux is not guaranteed to be a TensorTree and could contain strings etc.

The proposed fix is to return a new state with all other attributes modified in-place (i.e. pointers to old state) but aux replaced.

[KaelanDt]

Looks good to me, but in my opinion it would be good to check how this affects memory consumption and allocation - since you allocate a new state whenever you update it, which may be quite inefficient

[SamDuffield]

Looks good to me, but in my opinion it would be good to check how this affects memory consumption and allocation - since you allocate a new state whenever you update it, which may be quite inefficient

This PR shouldn't affect memory consumption, it just changes the handling of the algorithm states to a better convention.

It would be good to have some numerics on memory consumption and the pros and cons of using inplace and even whether we should continue to support it.

There is also an element of horses-for-courses since for MCMC-style where you want to collect samples along a trajectory you need inplace=False whereas for optimization or deep ensemble-style where you only care about the final result you may prefer inplace=True if its faster (which I'm not sure on).

[KaelanDt]

Looks good to me, but in my opinion it would be good to check how this affects memory consumption and allocation - since you allocate a new state whenever you update it, which may be quite inefficient

This PR shouldn't affect memory consumption, it just changes the handling of the algorithm states to a better convention.

It would be good to have some numerics on memory consumption and the pros and cons of using inplace and even whether we should continue to support it.

There is also an element of horses-for-courses since for MCMC-style where you want to collect samples along a trajectory you need inplace=False whereas for optimization or deep ensemble-style where you only care about the final result you may prefer inplace=True if its faster (which I'm not sure on).

I'm not sure about this: the previous inplace behaviour would change the elements of a previously allocated object. Now, you re-allocate a new object every time

[KaelanDt]

Looks good to me apart from a docstring to change

[SamDuffield]

Looks good to me, but in my opinion it would be good to check how this affects memory consumption and allocation - since you allocate a new state whenever you update it, which may be quite inefficient

This PR shouldn't affect memory consumption, it just changes the handling of the algorithm states to a better convention.
It would be good to have some numerics on memory consumption and the pros and cons of using inplace and even whether we should continue to support it.
There is also an element of horses-for-courses since for MCMC-style where you want to collect samples along a trajectory you need inplace=False whereas for optimization or deep ensemble-style where you only care about the final result you may prefer inplace=True if its faster (which I'm not sure on).

I'm not sure about this: the previous inplace behaviour would change the elements of a previously allocated object. Now, you re-allocate a new object every time

We define a new NamedTuple but if inplace=True all Tensors are pointers to the same memory as the previous ones (aside from aux which is not guaranteed to be a TensorTree). This is also checked in the tests.