[Bug] `I_like` does not copy the device of the LinearOperator it is "like"

[dashstander]

🐛 Bug

If I have a linear operator A on GPU and call Ident = cola.ops.I_like(A), then the new identity operator will be placed on CPU.

This is annoying because if combined functions like cola.eig(Ident - A, k=1) will fail because the the tensors are on different devices.

To reproduce

** Code snippet to reproduce **

import torch
import cola

print(f'CUDA Available {torch.cuda.is_available()}')

A = cola.ops.Diagonal(torch.randn((200, 200)))
A = A.to('cuda')

Ident = cola.ops.I_like(A)

print(f'A: {A.device}')
print(f'Ident: {Ident.device}')

** Stack trace/error message **

CUDA Available True
A: cuda:0
Ident: cpu

Expected Behavior

For I_like(A).device == A.device.

System information

Please complete the following information:
CoLA: 0.0.5
Pytorch: 2.3.0+cu121
Python: 3.11.9
Ubunutu 24.04

[AndPotap]

Thank you for pointing out this bug and thanks for the concise and well though code snippet to reproduce! I've just added a fix on #99. Also the problem of trying to evaluate cola.eig(Ident - A, k=1) was much deeper that what I initially thought, involving a fix not only to I_like but other ops.