Skip to content

Commit

Permalink
Merge branch 'main' into frighterafix#1505
Browse files Browse the repository at this point in the history
  • Loading branch information
ColCarroll authored Jun 6, 2022
2 parents 944af40 + 8e72c11 commit bd53773
Show file tree
Hide file tree
Showing 392 changed files with 10,778 additions and 4,681 deletions.
4 changes: 2 additions & 2 deletions discussion/examples/windowed_sampling.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -1837,9 +1837,9 @@
"WARNING:tensorflow:Note that RandomStandardNormal inside pfor op may not give same output as inside a sequential loop.\n",
"Fast window 75\n",
"Slow window 25\n",
"WARNING:tensorflow:5 out of the last 5 calls to \u003cfunction slow_window at 0x7f9456031ea0\u003e triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for more details.\n",
"WARNING:tensorflow:5 out of the last 5 calls to \u003cfunction slow_window at 0x7f9456031ea0\u003e triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for more details.\n",
"Slow window 50\n",
"WARNING:tensorflow:6 out of the last 6 calls to \u003cfunction slow_window at 0x7f9456031ea0\u003e triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for more details.\n",
"WARNING:tensorflow:6 out of the last 6 calls to \u003cfunction slow_window at 0x7f9456031ea0\u003e triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has reduce_retracing=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for more details.\n",
"Slow window 100\n",
"Slow window 200\n",
"Fast window 75\n",
Expand Down
6 changes: 6 additions & 0 deletions discussion/meads/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
# Tuning-Free Generalized Hamiltonian Monte Carlo

The notebook meads.ipynb (best run via Google Colab or Jupyter) implements and
demonstrates the Maximum-Eigenvalue Adaptation of Damping and Step size (MEADS)
algorithm from the paper "Tuning-Free Generalized Hamiltonian Monte Carlo" (to
appear at AISTATS 2022).
1,336 changes: 1,336 additions & 0 deletions discussion/meads/meads.ipynb

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ def _sample_posterior(target_log_prob_unconstrained,
parallel_iterations=10,
jit_compile=True,
use_input_signature=False,
experimental_relax_shapes=False):
reduce_retracing=False):
"""MCMC sampling with HMC/NUTS using an expanding epoch tuning scheme."""

seed_stream = tfp.util.SeedStream(seed, 'window_tune_nuts_sampling')
Expand Down Expand Up @@ -117,7 +117,7 @@ def _sample_posterior(target_log_prob_unconstrained,
input_signature=input_signature,
autograph=False,
jit_compile=jit_compile,
experimental_relax_shapes=experimental_relax_shapes)
reduce_retracing=reduce_retracing)
def fast_adaptation_interval(num_steps, previous_state):
"""Step size only adaptation interval.
Expand Down Expand Up @@ -179,7 +179,7 @@ def body_fn_window2(
input_signature=input_signature,
autograph=False,
jit_compile=jit_compile,
experimental_relax_shapes=experimental_relax_shapes)
reduce_retracing=reduce_retracing)
def slow_adaptation_interval(num_steps, previous_n, previous_state,
previous_mean, previous_cov):
"""Interval that tunes the mass matrix and step size simultaneously.
Expand Down Expand Up @@ -328,7 +328,7 @@ def window_tune_nuts_sampling(target_log_prob,
parallel_iterations=10,
jit_compile=True,
use_input_signature=True,
experimental_relax_shapes=False):
reduce_retracing=False):
"""Sample from a density with NUTS and an expanding window tuning scheme.
This function implements a turnkey MCMC sampling routine using NUTS and an
Expand All @@ -347,7 +347,7 @@ def window_tune_nuts_sampling(target_log_prob,
of the tuning epoch (window 1, 2, and 3 in Stan [1]) run with two @tf.function
compiled functions. The user can control the compilation options using the
kwargs `jit_compile`, `use_input_signature`, and
`experimental_relax_shapes`. Setting all to True would compile to XLA and
`reduce_retracing`. Setting all to True would compile to XLA and
potentially avoid the small overhead of function recompilation (note that it
is not yet the case in XLA right now). It is not yet clear whether doing it
this way is better than just wrapping the full inference routine in
Expand Down Expand Up @@ -403,7 +403,7 @@ def window_tune_nuts_sampling(target_log_prob,
function is always compiled by XLA.
use_input_signature: If True, generate an input_signature kwarg to pass to
tf.function decorator.
experimental_relax_shapes: kwarg pass to tf.function decorator. When True,
reduce_retracing: kwarg pass to tf.function decorator. When True,
tf.function may generate fewer, graphs that are less specialized on input
shapes.
Expand Down Expand Up @@ -564,6 +564,6 @@ def target_log_prob_unconstrained_concated(x):
parallel_iterations=parallel_iterations,
jit_compile=jit_compile,
use_input_signature=use_input_signature,
experimental_relax_shapes=experimental_relax_shapes)
reduce_retracing=reduce_retracing)
return forward_transform(
split_and_reshape(nuts_samples)), diagnostic, conditioning_bijector
6 changes: 3 additions & 3 deletions spinoffs/fun_mc/fun_mc/dynamic/backend_jax/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@

def map_tree(fn, tree, *args):
"""Maps `fn` over the leaves of a nested structure."""
return tree_util.tree_multimap(fn, tree, *args)
return tree_util.tree_map(fn, tree, *args)


def flatten_tree(tree):
Expand All @@ -66,7 +66,7 @@ def map_tree_up_to(shallow, fn, tree, *rest):
def wrapper(_, *rest):
return fn(*rest)

return tree_util.tree_multimap(wrapper, shallow, tree, *rest)
return tree_util.tree_map(wrapper, shallow, tree, *rest)


def get_shallow_tree(is_leaf, tree):
Expand All @@ -76,7 +76,7 @@ def get_shallow_tree(is_leaf, tree):

def assert_same_shallow_tree(shallow, tree):
"""Asserts that `tree` has the same shallow structure as `shallow`."""
# Do a dummy multimap for the side-effect of verifying that the structures are
# Do a dummy map for the side-effect of verifying that the structures are
# the same. This doesn't catch all the errors we actually care about, sadly.
map_tree_up_to(shallow, lambda *args: (), tree)

Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
1 change: 0 additions & 1 deletion spinoffs/inference_gym/inference_gym/backends/rewrite.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
1 change: 0 additions & 1 deletion spinoffs/inference_gym/inference_gym/backends/util.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
1 change: 0 additions & 1 deletion spinoffs/inference_gym/inference_gym/dynamic/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
1 change: 0 additions & 1 deletion spinoffs/inference_gym/inference_gym/internal/data.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
1 change: 0 additions & 1 deletion spinoffs/inference_gym/inference_gym/internal/data_test.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python2, python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2021 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2021 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
1 change: 0 additions & 1 deletion spinoffs/inference_gym/inference_gym/internal/test_util.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
1 change: 0 additions & 1 deletion spinoffs/inference_gym/inference_gym/targets/banana.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python2, python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python2, python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
# Lint as: python3
# Copyright 2020 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
Loading

0 comments on commit bd53773

Please sign in to comment.