Add kwargs to assert_jac_close and assert_grad_close

These functions really are wrappers around assert_close, so we'd like them to always also take the parameters of assert_close, even if those change in the future, and to have the same default values. So I think kwargs is justified here. Also it's not user facing so the lack of documentation of the expected types will not be visible.

Author: ValerianRey

tests/unit/autojac/_asserts.py

@@ -16,10 +16,10 @@ def assert_has_no_jac(t: torch.Tensor) -> None:
     assert not hasattr(t, "jac")
 
 
-def assert_jac_close(t: torch.Tensor, expected_jac: torch.Tensor) -> None:
+def assert_jac_close(t: torch.Tensor, expected_jac: torch.Tensor, **kwargs) -> None:
     assert hasattr(t, "jac")
     t_ = cast(TensorWithJac, t)
-    assert_close(t_.jac, expected_jac)
+    assert_close(t_.jac, expected_jac, **kwargs)
 
 
 def assert_has_grad(t: torch.Tensor) -> None:
@@ -30,6 +30,6 @@ def assert_has_no_grad(t: torch.Tensor) -> None:
     assert t.grad is None
 
 
-def assert_grad_close(t: torch.Tensor, expected_grad: torch.Tensor) -> None:
+def assert_grad_close(t: torch.Tensor, expected_grad: torch.Tensor, **kwargs) -> None:
     assert t.grad is not None
-    assert_close(t.grad, expected_grad)
+    assert_close(t.grad, expected_grad, **kwargs)