Support vmap over multiple instances of QumodeCircuit

[Jooyuza]

New features:

• set_fock_basis(): allow users manually set fock basis states before forward
• get_fock_basis(): check the output fock basis states in current settings
• optional bool arg sort of forward could decide whether to sort dictionary of Fock basis states in the descending order of probs. If sort=False, the dictionary will be in the lexicographic order.

Miscellaneous:

• update get_unitary: replacing in-place operations with out-of-place ops index_put to support vmap
• support return hash value of FockState on meta devices.
• fix typos

Example use case:

import torch
import torch.nn as nn

class BS(nn.Module):
    def __init__(self):
        super().__init__()
        cir = dq.QumodeCircuit(nmode=3, init_state=[1,1,0])
        cir.bs([0,1], encode=True)
        cir.ps([1], encode=True)
        cir.bs([1,2])
        cir.set_fock_basis() # set all fock basis before vmap
        self.cir = cir

    def forward(self, x):
        d = self.cir(x, is_prob=True, sort=False) # forward without extra sorting
        sorted_items = d.items()
        sorted_values = [v for k, v in sorted_items]
        probs = torch.cat(sorted_values,dim=1)
        return probs

num_models = 5
batch_size = 64

models = [BS() for i in range(num_models)]
data = torch.randn(batch_size, 3)

def wrapper(params, buffers, data):
    return torch.func.functional_call(models[0], (params, buffers), data)

params, buffers = torch.func.stack_module_state(models)
output = vmap(wrapper, (0, 0, None))(params, buffers, data)

print(output.shape)

[sansiro77]

LGTM