What does this PR do?

First pass implementation of KPlanes implementation (https://brentyi.github.io/tilted/, https://sarafridov.github.io/K-Planes/). These tensor decomposition models require multiple parmeter optimization which is not handled in OptimizerMixin.

The solution @arthurmccray and I came up with is to setup a per-parameter learning rate PPLR abstract class which will be instantiated with models that have multiple parameters to be optimized to. This is implemented incore/ml/models/model_base.py. The most important part about this class is the get_params method, which is used for parsing the trainable parameters (see ObjectTensorDecomp's get_optimization_parameters) .

class PPLR(ABC):
    """
    Abstract base class for models that require multi-scale parameter optimization.
    """

    @abstractmethod
    def get_params(self) -> Dict[str, list[nn.Parameter]]:
        """
        Return a dictionary of parameters grouped by key.

        For example if your nn.Module has multiple optimizable parameter groups,
        you can return a dictionary with the keys "grids" and "sigma_net"
        (KPlanes example).
        """
        pass

The KPlanes and TILTED implementations are in core/ml/models in kplanes.py with KPlanes, KPlanesTILTED, and CPTilted. Note CPTIlted is used for the two-phase warmup that the paper recommends by pretraining the SO3 rotations of the planes in a lower representation space. so3params.py also has both quaternion and R9+SVD implementations which can be swapped as a parameter in KPlanesTILTED.

In tomography/object_models.py there is now a new object that handles the new tensor decomposition methods, ObjectTensorDecomp that inherits from ObjectINR and the following overloads are needed:

• optimizer_params property: The setter is changed, since we're building a a dictionary of parameters into our optimizers per https://docs.pytorch.org/docs/stable/optim.html#per-parameter-options.
• reconnect_optimizer_to_parameters: This one is the most annoying thing to implement, since now our parameters is a dictionary.

There are a few optional overrides that I've also included preemptively just incase if I want to add additional soft/hard constraints to these methods as you can see both apply_soft_constraints and apply_hard_constraints are still there.

Key change to OptimizerMixin: I have changed set_optimizer to handle a dictionary of parameters. I think this is better than my original idea of just implementing this in object_models.py since you would have to overload that function again. I think this is fine and still works for single optimization stuff.

There are also small changes to different aspects of the Tomography module, but mostly contained in object_models.py. Listing out the changes and rationale in the tomography files:

• tomography.py: I've disabled torch.bfloat16 autocasting since F.grid_sample does not have autograd implemented errors. I've also put the TILTED tensor decomp pretraining check by looking at the convergence of the SO3 rotations. I don't like this being here, but I'm not quite sure where else to put this.
• tomography_base.py: Type-hinting change for setting up distributed.
• tomography_opt.py: Parsing is already handled in object_models.py.

What the reviewer should do

Attached is a notebook that tests tensor decomposition methods in tomography: 0415_PPLR_testing.ipynb. The dataset here is from quantem-tutorials, so just set the directory of the same dataset to this notebook. The runtime can be increased by increasing the batch size, but the learning rates have to be scaled. Setting the batch size to 4096 seems to be fine? So feel free to change it and scale the learning rates by 2.

• Double check implementations across the board, especially the optimizer_mixin.py, and object_models.py
• Organization makes sense for tensor decomposition stuff. I'm actually not very happy with how core/ml is looking. Where do we draw the line of what lives in models (I think everything model related should live in there tbh: INR, CNN, etc...)?
• Check reconstruction workflows are still working i.e, ptychography due to the OptimizerMixin change.