From 45c47e657b0a674f5a3ed4a3b2c2ca13b811f194 Mon Sep 17 00:00:00 2001
From: Faizan Muhammad <fmuham@google.com>
Date: Wed, 2 Nov 2022 13:49:33 -0700
Subject: [PATCH] Remove functools.wraps for unbound methods to support
 inspect.signature

PiperOrigin-RevId: 485685579
---
 sonnet/src/utils.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/sonnet/src/utils.py b/sonnet/src/utils.py
index 5c37fb9a..07a9b7f8 100644
--- a/sonnet/src/utils.py
+++ b/sonnet/src/utils.py
@@ -77,7 +77,9 @@ def _decorate_bound_method(*args, **kwargs):
     argspec = inspect.getfullargspec(f)
     if argspec.args and argspec.args[0] == "self":
 
-      @functools.wraps(f)
+      # TODO(b/123870292): Bound methods with multiple decorators can be
+      # detected as unbound too, which breaks tf.function usage.
+      # Attempt migrating to wrapt instead of using this custom approach.
       def _decorate_unbound_method(self, *args, **kwargs):
         bound_method = f.__get__(self, self.__class__)  # pytype: disable=attribute-error
         return decorator_fn(bound_method, self, args, kwargs)