torch.sparse.as_sparse_gradcheck

torch.sparse.as_sparse_gradcheck(gradcheck) [исходный код]

Декоратор функции, расширяющий gradcheck для разреженных тензоров.

Декоратор для torch.autograd.gradcheck или его вариантов functools.partial, расширяющий функцию gradcheck поддержкой входных функций, которые работают с разреженными тензорами и/или возвращают их.

Гарантируется, что указанная функция gradcheck работает только с тензорами со страйдовым хранением.

Например:

>>> warnings.filterwarnings(
...     "ignore", message=".*Sparse CSR tensor support is in beta state"
... )  # docs: hide
>>> gradcheck = torch.sparse.as_sparse_gradcheck(torch.autograd.gradcheck)
>>> x = (
...     torch.tensor([[0, 1], [2, 3]], dtype=torch.float64)
...     .to_sparse_coo()
...     .requires_grad_(True)
... )
>>> gradcheck(lambda x: x.to_sparse_csr(), x)
True