[UnitaryHack] Make ansatz class a training hyperparameter

[neiljdo]

PR adds the ansatz class and necessary arguments for instantiation (e.g. ob_map and ansatz-specific kwargs) as hyperparameters to the trainer class. The functions related to creating and loading checkpoints have also been updated to accommodate the new ansatz class hyperparameter.

Addresses https://github.com/CQCL/lambeq/issues/86.

[ianyfan]

Hi Neil, thanks for your submission, it looks quite good already. I've enabled the tests to run on this PR, so would you be able to have a look at where they're failing and see if you could fix them up?

[neiljdo]

Hi @ianyfan, thank you for the review. I just pushed a new commit to address the flake8 violations. I'll wait for the workflow run results to see if I missed some.

[neiljdo]

@ianyfan just saw the failing jobs related to the notebook tests and type checking - will push a fix to those soon.

[dimkart]

@ianyfan just saw the failing jobs related to the notebook tests and type checking - will push a fix to those soon.

@neiljdo Hi, it would be great if you could that by tomorrow, 13/6, which is the last day of the hackathon.

[nathanshammah]

@dimkart @neiljdo PRs can be reviewed also shortly after June 13th and still be awarded a bounty in uHACK if accepted.

[dimkart]

@nathanshammah

PRs can be reviewed also shortly after June 13th and still be awarded a bounty in uHACK if accepted.

Thanks, noted.

[nikhilkhatri]

Would also be great if you could run the clean_notebooks.py script, which removes some metadata from the examples and tutorials.

[neiljdo]

@nikhilkhatri, Thank you for the review. Here are my thoughts on your comments above.

1. We should save the circuits in order to do the diagram-to-circuit conversion only once. I'll implement a fix for this one right away.
2. For the ansatz serialization, I only save the components (i.e. the ansatz class, ob_map, and other ansatz-specific kwargs) required to initialize the ansatz during trainer checkpoint creation and re-initialize the ansatz during checkpoint load. I thought the ansatz deterministic and could be re-instantiated, so I went this route. I would like to do tests as part of the testing suite (not manually) - this implies adding new tests that use different ansatzes for the trainers.

For (3) and (4) above, I realized that the model is coupled with the circuits while working on the issue. I saw this more while updating the demo notebooks and saw that the Model.from_diagram call also needs to get the test circuits. Let's discuss these two later once I've resolved the first two.

[nikhilkhatri]

@neiljdo you are correct in that the ansatze are deterministic. My point was to ensure that there are no Exceptions raised when pickling an ansatz class (for example, if they use a lambda function).

[nikhilkhatri]

w.r.t the lambda pickling, your code may be alright since it pickles the class, not the instance. It would still be good to make sure.

[ACE07-Sev]

May I drop a suggestion here? Based on what I understand from the code :

1) If Model is from checkpoint, so you just skip the _init_model_from_datasets method , but still have to recreate the circuits at training_step and validation_step (assuming number 1 is fixed).

2) If Model is not from checkpoint, you make the circuits at both the _init_model_from_datasets method and the training_step and validation_step.

Since the dataset method is not run every time, I think it could be better if you make the circuits once, pass the list to _init_model_from_datasets method, and also use the list in the training and validation step (based on the index of the diagram of course).

My suggestion is after fixing 1, to have the _init_model_from_datasets also have the reference from 1's fix. @nikhilkhatri dear, am I correct? Please correct me if I am wrong.

[nikhilkhatri]

Hi @ACE07-Sev , Thanks for your comments. Indeed, the circuits should only be created once, and then used wherever the ansatz is currently applied. The circuits need to be created wherever the dataset and ansatz parameters are first available. If the dataset is taken as input first during fit, then this is where the ansatz can be applied, and circuits from this can be reused throughout the setup, training and validation.

[neiljdo]

Hi @nikhilkhatri,

I've addressed points (1) and (2) from your original comment in my latest update. Specifically,

1. Circuits are now generated once at the start of the Trainer.fit() execution. Circuits are also included in the checkpoint.
2. I've used dill as a drop-in replacement for pickle for checkpoint generation as it handles lambda functions well. I've added tests for all available ansatzes to make sure that it still generates the correct output after serialization-deserialization. With this, we can simplify the ansatz checkpoint creation i.e. no need to re-instantiate.

Let me know if there are still issues, thank you very much!

EDIT: I ran the clean_notebooks.py script but it did not update the notebooks.

[nikhilkhatri]

Hi @neiljdo Thanks a lot for these changes!

I have a couple of comments:

1. Saving all the circuits will create a very large checkpoint file, depending on the dataset size, and circuit complexity. I dont think this is required. If a trainer is loaded from a checkpoint, the user will likely call fit, at which point they can provide the necessary circuits.
2. The tests seem to be failing because dill hasnt been added to requirements.
3. (Minor) The tests for pickling the ansatze are very verbose. Using something like this can help here.

It would be great if you could address issues 1 and 2 for the hack.

[neiljdo]

@nikhilkhatri Noted, thanks. I'll remove the circuits from the checkpoint.

[ACE07-Sev]

Neil dear :

lambeq/training/checkpoint.py:27: error: Cannot find implementation or library stub for module named "dill" [import]

The test with pytest passed, so I think you just need to fix this import error.

[dimkart]

@neiljdo Hey Neil please leave a comment in the Issue #86, so you can be assigned to it (otherwise your handle cannot be found).

[dimkart]

Moved to lambeq's private repository for further development. This feature will be included in one of the upcoming releases.