Improve efficiency and fix bugs (device and dtype) for Gaussian backend

[Jooyuza]

No description provided.

[sansiro77]

It seems to be a bug related to vmap + slice assignment + cuda.

For example,

x = torch.randn(2, 4, 4, device='cuda:0')
def f(x):
    x[0:2, 0:2] = 1
    return x
torch.vmap(f)(x)

is OK.

However,

x = torch.randn(2, 4, 4, device='cuda:0')
def f(x):
    x[[0,1], [0,1]] = 1
    return x
torch.vmap(f)(x)

reports error.

Try something like:

mat = torch.randn(3, 3, device='cuda:0')
x = mat.new_ones(1)
torch.block_diag(*([x] * 3))

[sansiro77]

Try

mat = torch.randn(3, 3, device='cuda:0')
x = mat.new_ones(3)
torch.diag(x)

to improve efficiency.

[Jooyuza]

[sansiro77]

Good job!