Switched from the deprecated defun_with_attributes() to tf.function()

PiperOrigin-RevId: 511782775
google-deepmind · Feb 23, 2023 · 26a685e · 26a685e
1 parent ee36d84
commit 26a685e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/requirements-tf.txt b/requirements-tf.txt
@@ -1,2 +1,2 @@
-tensorflow==2.11.0
+tensorflow==2.12.0rc0
 tensorflow-probability==0.12.2
diff --git a/sonnet/src/recurrent.py b/sonnet/src/recurrent.py
@@ -34,7 +34,6 @@
 # pylint: disable=g-direct-tensorflow-import
 # Required for specializing `UnrolledLSTM` per device.
 from tensorflow.python import context as context_lib
-from tensorflow.python.eager import function as function_lib
 # pylint: enable=g-direct-tensorflow-import
 
 
@@ -1029,9 +1028,9 @@ def wrapper(*args, **kwargs):
     unique_api_name = "{}_{}".format(api_name, uuid.uuid4())
     functions = {}
     for device, specialization in specializations.items():
-      functions[device] = function_lib.defun_with_attributes(
+      functions[device] = tf.function(
           specialization,
-          attributes={
+          experimental_attributes={
               "api_implements": unique_api_name,
               "api_preferred_device": device
           })