# This file is part of tad-mctc.
#
# SPDX-Identifier: Apache-2.0
# Copyright (C) 2024 Grimme Group
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Autograd Utility: Gradcheck
===========================
Collection of utility functions for testing.
"""
from __future__ import annotations
from torch.autograd.gradcheck import gradcheck, gradgradcheck
from ..typing import Any, Callable, Protocol, Tensor, TensorOrTensors
__all__ = ["dgradcheck", "dgradgradcheck"]
FAST_MODE = True
"""Default for fast_mode argument (True)."""
class _GradcheckFunction(Protocol):
"""
Type annotation for gradcheck function.
"""
def __call__( # type: ignore
self,
func: Callable[..., TensorOrTensors],
inputs: TensorOrTensors,
*,
eps: float = 1e-6,
atol: float = 1e-5,
rtol: float = 1e-3,
raise_exception: bool = True,
check_sparse_nnz: bool = False,
nondet_tol: float = 0.0,
check_undefined_grad: bool = True,
check_grad_dtypes: bool = False,
check_batched_grad: bool = False,
check_batched_forward_grad: bool = False,
check_forward_ad: bool = False,
check_backward_ad: bool = True,
fast_mode: bool = False,
) -> bool: ...
class _GradgradcheckFunction(Protocol):
"""
Type annotation for gradgradcheck function.
"""
def __call__( # type: ignore
self,
func: Callable[..., TensorOrTensors],
inputs: TensorOrTensors,
grad_outputs: TensorOrTensors | None = None,
*,
eps: float = 1e-6,
atol: float = 1e-5,
rtol: float = 1e-3,
gen_non_contig_grad_outputs: bool = False,
raise_exception: bool = True,
nondet_tol: float = 0.0,
check_undefined_grad: bool = True,
check_grad_dtypes: bool = False,
check_batched_grad: bool = False,
check_fwd_over_rev: bool = False,
check_rev_over_rev: bool = True,
fast_mode: bool = False,
) -> bool: ...
def _wrap_gradcheck(
gradcheck_func: _GradcheckFunction | _GradgradcheckFunction,
func: Callable[..., TensorOrTensors],
diffvars: TensorOrTensors,
**kwargs: Any,
) -> bool:
fast_mode = kwargs.pop("fast_mode", FAST_MODE)
try:
assert gradcheck_func(func, diffvars, fast_mode=fast_mode, **kwargs)
finally:
if isinstance(diffvars, Tensor):
diffvars.detach_()
elif isinstance(diffvars, (list, tuple)):
for diffvar in diffvars:
if isinstance(diffvar, Tensor):
diffvar.detach_()
return True
[docs]
def dgradcheck(
func: Callable[..., TensorOrTensors],
diffvars: TensorOrTensors,
**kwargs: Any,
) -> bool:
"""
Wrapper for `torch.autograd.gradcheck` that detaches the differentiated
variables after the check.
Parameters
----------
func : Callable[..., TensorOrTensors]
Forward function.
diffvars : TensorOrTensors
Variables w.r.t. which we differentiate.
Returns
-------
bool
Status of check.
"""
return _wrap_gradcheck(gradcheck, func, diffvars, **kwargs)
[docs]
def dgradgradcheck(
func: Callable[..., TensorOrTensors],
diffvars: TensorOrTensors,
**kwargs: Any,
) -> bool:
"""
Wrapper for `torch.autograd.gradgradcheck` that detaches the differentiated
variables after the check.
Parameters
----------
func : Callable[..., TensorOrTensors]
Forward function.
diffvars : TensorOrTensors
Variables w.r.t. which we differentiate.
Returns
-------
bool
Status of check.
"""
return _wrap_gradcheck(gradgradcheck, func, diffvars, **kwargs)