Add a prototype functional MCMC API

PiperOrigin-RevId: 242552429

experimental/fun_mcmc/BUILD

@@ -0,0 +1,63 @@
+# Copyright 2018 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:
+#   Functional MCMC API.
+
+# [internal] load pytype.bzl
+
+licenses(["notice"])  # Apache 2.0
+
+package(
+    default_visibility = [
+        "//tensorflow_probability:__subpackages__",
+    ],
+)
+
+exports_files(["LICENSE"])
+
+# pytype
+py_library(
+    name = "fun_mcmc",
+    srcs = ["__init__.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":fun_mcmc_lib",
+    ],
+)
+
+# pytype
+py_library(
+    name = "fun_mcmc_lib",
+    srcs = ["fun_mcmc_lib.py"],
+    srcs_version = "PY2AND3",
+    deps = [
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+    ],
+)
+
+py_test(
+    name = "fun_mcmc_test",
+    srcs = ["fun_mcmc_test.py"],
+    deps = [
+        ":fun_mcmc",
+        # absl/testing:parameterized dep,
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)

experimental/fun_mcmc/__init__.py

@@ -0,0 +1,55 @@
+# Copyright 2018 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.
+# ============================================================================
+"""Functional MCMC."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from experimental.fun_mcmc.fun_mcmc_lib import call_and_grads
+from experimental.fun_mcmc.fun_mcmc_lib import call_fn
+from experimental.fun_mcmc.fun_mcmc_lib import hamiltonian_monte_carlo
+from experimental.fun_mcmc.fun_mcmc_lib import HamiltonianMonteCarloExtra
+from experimental.fun_mcmc.fun_mcmc_lib import HamiltonianMonteCarloState
+from experimental.fun_mcmc.fun_mcmc_lib import leapfrog_step
+from experimental.fun_mcmc.fun_mcmc_lib import LeapFrogStepExtras
+from experimental.fun_mcmc.fun_mcmc_lib import LeapFrogStepState
+from experimental.fun_mcmc.fun_mcmc_lib import maybe_broadcast_structure
+from experimental.fun_mcmc.fun_mcmc_lib import metropolis_hastings_step
+from experimental.fun_mcmc.fun_mcmc_lib import PotentialFn
+from experimental.fun_mcmc.fun_mcmc_lib import sign_adaptation
+from experimental.fun_mcmc.fun_mcmc_lib import State
+from experimental.fun_mcmc.fun_mcmc_lib import trace
+from experimental.fun_mcmc.fun_mcmc_lib import transform_log_prob_fn
+from experimental.fun_mcmc.fun_mcmc_lib import TransitionOperator
+
+__all__ = [
+    'HamiltonianMonteCarloExtra',
+    'HamiltonianMonteCarloState',
+    'LeapFrogStepExtras',
+    'LeapFrogStepState',
+    'PotentialFn',
+    'State',
+    'TransitionOperator',
+    'call_and_grads',
+    'call_fn',
+    'hamiltonian_monte_carlo',
+    'maybe_broadcast_structure',
+    'metropolis_hastings_step',
+    'leapfrog_step',
+    'sign_adaptation',
+    'trace',
+    'transform_log_prob_fn',
+]