Add experimental joint-distribution layers library.

PiperOrigin-RevId: 396021661

Author: SiegeLordEx

tensorflow_probability/python/build_defs.bzl

@@ -18,7 +18,7 @@ NO_REWRITE_NEEDED = [
     "internal:all_util",
     "internal:docstring_util",
     "internal:reparameterization",
-    "layers",
+    "python/layers",
     "platform_google",
 ]
 

tensorflow_probability/python/experimental/BUILD

@@ -53,6 +53,7 @@ multi_substrate_py_library(
         "//tensorflow_probability/python/experimental/bijectors",
         "//tensorflow_probability/python/experimental/distribute",
         "//tensorflow_probability/python/experimental/distributions",
+        "//tensorflow_probability/python/experimental/joint_distribution_layers",
         "//tensorflow_probability/python/experimental/linalg",
         "//tensorflow_probability/python/experimental/marginalize",
         "//tensorflow_probability/python/experimental/math",

tensorflow_probability/python/experimental/__init__.py

@@ -35,6 +35,7 @@
 from tensorflow_probability.python.experimental import bijectors
 from tensorflow_probability.python.experimental import distribute
 from tensorflow_probability.python.experimental import distributions
+from tensorflow_probability.python.experimental import joint_distribution_layers
 from tensorflow_probability.python.experimental import linalg
 from tensorflow_probability.python.experimental import marginalize
 from tensorflow_probability.python.experimental import math
@@ -62,6 +63,7 @@
     'bijectors',
     'distribute',
     'distributions',
+    'joint_distribution_layers',
     'linalg',
     'marginalize',
     'math',

tensorflow_probability/python/experimental/joint_distribution_layers/BUILD

@@ -0,0 +1,74 @@
+# 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.
+# ============================================================================
+# Description:
+#   Joint Distribution Layers library.
+
+load(
+    "//tensorflow_probability/python:build_defs.bzl",
+    "multi_substrate_py_library",
+    "multi_substrate_py_test",
+)
+
+package(
+    default_visibility = [
+        "//tensorflow_probability:__subpackages__",
+    ],
+)
+
+licenses(["notice"])
+
+multi_substrate_py_library(
+    name = "joint_distribution_layers",
+    srcs = ["__init__.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":layers",
+    ],
+)
+
+multi_substrate_py_library(
+    name = "layers",
+    srcs = ["layers.py"],
+    srcs_version = "PY3",
+    deps = [
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/distributions:deterministic",
+        "//tensorflow_probability/python/distributions:joint_distribution_coroutine",
+        "//tensorflow_probability/python/distributions:normal",
+        "//tensorflow_probability/python/distributions:sample",
+        "//tensorflow_probability/python/internal:vectorization_util",
+    ],
+)
+
+multi_substrate_py_test(
+    name = "layers_test",
+    size = "medium",
+    srcs = ["layers_test.py"],
+    srcs_version = "PY3",
+    deps = [
+        # absl/testing:parameterized dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:custom_gradient",
+        "//tensorflow_probability/python/internal:prefer_static",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)
+
+exports_files(
+    glob(["**/*.py"]),
+    visibility = ["//tensorflow_probability:__subpackages__"],
+)

tensorflow_probability/python/experimental/joint_distribution_layers/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+# ============================================================================
+"""Experimental Joint Distribution Layers library."""
+
+from tensorflow_probability.python.experimental.joint_distribution_layers import layers
+from tensorflow_probability.python.experimental.joint_distribution_layers.layers import *
+
+__all__ = layers.__all__