Update BfgsTest and LBfgsTest to work with tf.function layering effor…

…ts by removing wrapping from val_and_grad. PiperOrigin-RevId: 498781838
tensorflow · Dec 31, 2022 · 5ce266f · 5ce266f
1 parent 517027b
commit 5ce266f
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 6 deletions.
diff --git a/tensorflow_probability/python/optimizer/bfgs_test.py b/tensorflow_probability/python/optimizer/bfgs_test.py
@@ -14,8 +14,6 @@
 # ============================================================================
 """Tests for the unconstrained BFGS optimizer."""
 
-import functools
-
 from absl.testing import parameterized
 import numpy as np
 from scipy.stats import special_ortho_group
@@ -30,7 +28,6 @@
 
 
 def _make_val_and_grad_fn(value_fn):
-  @functools.wraps(value_fn)
   def val_and_grad(x):
     return gradient.value_and_gradient(value_fn, x)
   return val_and_grad

diff --git a/tensorflow_probability/python/optimizer/lbfgs_test.py b/tensorflow_probability/python/optimizer/lbfgs_test.py
@@ -14,8 +14,6 @@
 # ============================================================================
 """Tests for the unconstrained L-BFGS optimizer."""
 
-import functools
-
 from absl.testing import parameterized
 import numpy as np
 from scipy.stats import special_ortho_group
@@ -30,7 +28,6 @@
 
 
 def _make_val_and_grad_fn(value_fn):
-  @functools.wraps(value_fn)
   def val_and_grad(x):
     return gradient.value_and_gradient(value_fn, x)
   return val_and_grad