Patch Numpy and JAX backends before converting DeferredTensor and TransformedVariable to CompositeTensor.

emilyfertig · brianwa84 · commit 673ff8924570 · 2021-05-20T13:45:39.000-04:00
PiperOrigin-RevId: 373906496
diff --git a/tensorflow_probability/python/internal/backend/jax/BUILD b/tensorflow_probability/python/internal/backend/jax/BUILD
@@ -53,6 +53,7 @@ FILENAMES = [
     "private",
     "random_generators",
     "raw_ops",
+    "resource_variable_ops",
     "sets_lib",
     "sparse_lib",
     "tensor_array_ops",
diff --git a/tensorflow_probability/python/internal/backend/numpy/BUILD b/tensorflow_probability/python/internal/backend/numpy/BUILD
@@ -52,6 +52,7 @@ py_library(
         ":private",
         ":random_generators",
         ":raw_ops",
+        ":resource_variable_ops",
         ":sets_lib",
         ":sparse_lib",
         ":static_rewrites",
@@ -332,6 +333,11 @@ py_library(
     ],
 )
 
+py_library(
+    name = "resource_variable_ops",
+    srcs = ["resource_variable_ops.py"],
+)
+
 py_library(
     name = "sets_lib",
     srcs = ["sets_lib.py"],
diff --git a/tensorflow_probability/python/internal/backend/numpy/__init__.py b/tensorflow_probability/python/internal/backend/numpy/__init__.py
@@ -46,6 +46,8 @@
 from tensorflow_probability.python.internal.backend.numpy.numpy_array import *  # pylint: disable=wildcard-import
 from tensorflow_probability.python.internal.backend.numpy.numpy_math import *  # pylint: disable=wildcard-import
 from tensorflow_probability.python.internal.backend.numpy.ops import *  # pylint: disable=wildcard-import
+from tensorflow_probability.python.internal.backend.numpy.type_spec import BatchableTypeSpec
+from tensorflow_probability.python.internal.backend.numpy.type_spec import TypeSpec
 
 
 Assert = debugging.Assert
diff --git a/tensorflow_probability/python/internal/backend/numpy/ops.py b/tensorflow_probability/python/internal/backend/numpy/ops.py
@@ -61,9 +61,7 @@
     'Module',
     'Tensor',
     'TensorSpec',
-    'TypeSpec',
     'Variable',
-    'VariableSpec',
     # 'gradients',
 ]
 
@@ -696,15 +694,10 @@ class Tensor(six.with_metaclass(_TensorMeta)):
 
 
 class TensorSpec(object):
-  pass
 
-
-class TypeSpec(object):
-  pass
-
-
-class VariableSpec(object):
-  pass
+  def __init__(self, *args, **kwargs):
+    del args, kwargs
+    self.dtype = None
 
 
 class Module(object):
diff --git a/tensorflow_probability/python/internal/backend/numpy/resource_variable_ops.py b/tensorflow_probability/python/internal/backend/numpy/resource_variable_ops.py
@@ -0,0 +1,26 @@
+# Copyright 2021 The TensorFlow Probability Authors.
+#
+# 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.
+# ============================================================================
+"""Numpy stub for `resource_variable_ops`."""
+
+__all__ = [
+    'VariableSpec',
+]
+
+
+class VariableSpec(object):
+
+  def __init__(self, *args, **kwargs):
+    del args, kwargs
+    self.dtype = None
diff --git a/tensorflow_probability/python/internal/backend/numpy/tf_inspect.py b/tensorflow_probability/python/internal/backend/numpy/tf_inspect.py
@@ -27,6 +27,7 @@
 # Although `inspect` is different between Python 2 and 3, we should only ever
 # be using Python 3's inspect because JAX is Python 3 only and if TF is present
 # we will use `tf_inspect` which is compatible with both Python 2 and 3.
+Parameter = inspect.Parameter
 getfullargspec = inspect.getfullargspec
 getcallargs = inspect.getcallargs
 getframeinfo = inspect.getframeinfo
@@ -46,4 +47,5 @@
 ismethod = inspect.ismethod
 ismodule = inspect.ismodule
 isroutine = inspect.isroutine
+signature = inspect.signature
 stack = inspect.stack
diff --git a/tensorflow_probability/python/internal/backend/numpy/type_spec.py b/tensorflow_probability/python/internal/backend/numpy/type_spec.py
@@ -16,7 +16,9 @@
 
 __all__ = [
     'lookup',
-    'register'
+    'register',
+    'BatchableTypeSpec',
+    'TypeSpec',
 ]
 
 
@@ -30,3 +32,11 @@ def decorator_fn(cls):
 def lookup(_):
   # Raise ValueError instead of NotImplementedError to conform to TF.
   raise ValueError('`TypeSpec`s are not registered in Numpy/JAX.')
+
+
+class TypeSpec(object):
+  pass
+
+
+class BatchableTypeSpec(TypeSpec):
+  pass
diff --git a/tensorflow_probability/python/internal/backend/numpy/v2.py b/tensorflow_probability/python/internal/backend/numpy/v2.py
@@ -51,6 +51,8 @@
 from tensorflow_probability.python.internal.backend.numpy.numpy_math import *  # pylint: disable=wildcard-import
 from tensorflow_probability.python.internal.backend.numpy.ops import *  # pylint: disable=wildcard-import
 from tensorflow_probability.python.internal.backend.numpy.tensor_array_ops import TensorArray
+from tensorflow_probability.python.internal.backend.numpy.type_spec import BatchableTypeSpec
+from tensorflow_probability.python.internal.backend.numpy.type_spec import TypeSpec
 # pylint: enable=unused-import
 
 
diff --git a/tensorflow_probability/substrates/meta/rewrite.py b/tensorflow_probability/substrates/meta/rewrite.py
@@ -61,7 +61,7 @@
     ('from tensorflow.python.ops import '
      'resource_variable_ops'):
         ('from tensorflow_probability.python.internal.backend.numpy '
-         'import ops'),
+         'import resource_variable_ops'),
     'from tensorflow.python.util import':
         'from tensorflow_probability.python.internal.backend.numpy import',
     'from tensorflow.python.util.all_util':
