Probability: Is it possible to get reproducible samples from tfp.mcmc.sample_chain?

[aidinhass]

Running a simple Bayesian regression model, I am not able to replicate the results with multiple runs on GPU. I am wondering how I can set tfp.mcmc.sample_chain to generate reproducible results on GPU? Seeding the sample_chain didn't work for me.

import os
import random
from pprint import pprint
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import pandas as pd

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()

import tensorflow_probability as tfp

sns.reset_defaults()
#sns.set_style('whitegrid')
#sns.set_context('talk')
sns.set_context(context='talk',font_scale=0.7)

%config InlineBackend.figure_format = 'retina'
%matplotlib inline

tfd = tfp.distributions
tfb = tfp.bijectors

dtype = tf.float64
dfhogg = pd.DataFrame(np.array([[1, 201, 592, 61, 9, -0.84],
                                 [2, 244, 401, 25, 4, 0.31],
                                 [3, 47, 583, 38, 11, 0.64],
                                 [4, 287, 402, 15, 7, -0.27],
                                 [5, 203, 495, 21, 5, -0.33],
                                 [6, 58, 173, 15, 9, 0.67],
                                 [7, 210, 479, 27, 4, -0.02],
                                 [8, 202, 504, 14, 4, -0.05],
                                 [9, 198, 510, 30, 11, -0.84],
                                 [10, 158, 416, 16, 7, -0.69],
                                 [11, 165, 393, 14, 5, 0.30],
                                 [12, 201, 442, 25, 5, -0.46],
                                 [13, 157, 317, 52, 5, -0.03],
                                 [14, 131, 311, 16, 6, 0.50],
                                 [15, 166, 400, 34, 6, 0.73],
                                 [16, 160, 337, 31, 5, -0.52],
                                 [17, 186, 423, 42, 9, 0.90],
                                 [18, 125, 334, 26, 8, 0.40],
                                 [19, 218, 533, 16, 6, -0.78],
                                 [20, 146, 344, 22, 5, -0.56]]),
                   columns=['id','x','y','sigma_y','sigma_x','rho_xy'])


## for convenience zero-base the 'id' and use as index
dfhogg['id'] = dfhogg['id'] - 1
dfhogg.set_index('id', inplace=True)

## standardize (mean center and divide by 1 sd)
dfhoggs = (dfhogg[['x','y']] - dfhogg[['x','y']].mean(0)) / dfhogg[['x','y']].std(0)
dfhoggs['sigma_y'] = dfhogg['sigma_y'] / dfhogg['y'].std(0)
dfhoggs['sigma_x'] = dfhogg['sigma_x'] / dfhogg['x'].std(0)

X_np = dfhoggs['x'].values
sigma_y_np = dfhoggs['sigma_y'].values
Y_np = dfhoggs['y'].values

def sample(seed):

  mdl_ols_batch = tfd.JointDistributionSequential([
      # b0
      tfd.Normal(loc=tf.cast(0, dtype), scale=1.),
      # b1
      tfd.Normal(loc=tf.cast(0, dtype), scale=1.),
      # likelihood
      #   Using Independent to ensure the log_prob is not incorrectly broadcasted
      lambda b1, b0: tfd.Independent(
          tfd.Normal(
              # Parameter transformation
              loc=b0[..., tf.newaxis] + b1[..., tf.newaxis]*X_np[tf.newaxis, ...],
              scale=sigma_y_np[tf.newaxis, ...]),
          reinterpreted_batch_ndims=1
      ),
  ])


  @tf.function(autograph=False, experimental_compile=True)
  def run_chain(init_state, 
                step_size,
                target_log_prob_fn,
                unconstraining_bijectors,
                num_steps=500, 
                burnin=50):

    def trace_fn(_, pkr):
      return (
          pkr.inner_results.inner_results.target_log_prob,
          pkr.inner_results.inner_results.leapfrogs_taken,
          pkr.inner_results.inner_results.has_divergence,
          pkr.inner_results.inner_results.energy,
          pkr.inner_results.inner_results.log_accept_ratio
      )

    kernel = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=tfp.mcmc.NoUTurnSampler(
        target_log_prob_fn,
        step_size=step_size),
      bijector=unconstraining_bijectors)

    hmc = tfp.mcmc.DualAveragingStepSizeAdaptation(
      inner_kernel=kernel,
      num_adaptation_steps=burnin,
      step_size_setter_fn=lambda pkr, new_step_size: pkr._replace(
          inner_results=pkr.inner_results._replace(step_size=new_step_size)),
      step_size_getter_fn=lambda pkr: pkr.inner_results.step_size,
      log_accept_prob_getter_fn=lambda pkr: pkr.inner_results.log_accept_ratio
    )

    # Sampling from the chain.
    chain_state, sampler_stat = tfp.mcmc.sample_chain(
        num_results=num_steps,
        num_burnin_steps=burnin,
        current_state=init_state,
        kernel=hmc,
        trace_fn=trace_fn,
        seed=seed
    )
    return chain_state, sampler_stat

  nchain = 4
  b0, b1, _ = mdl_ols_batch.sample(nchain)
  init_state = [b0, b1]
  step_size = [tf.cast(i, dtype=dtype) for i in [.1, .1]]
  target_log_prob_fn = lambda *x: mdl_ols_batch.log_prob(x + (Y_np, ))

  # bijector to map contrained parameters to real
  unconstraining_bijectors = [
      tfb.Identity(),
      tfb.Identity(),
  ]

  samples, sampler_stat = run_chain(
      init_state, step_size, target_log_prob_fn, unconstraining_bijectors)
  print(tf.reduce_sum(samples))


seed = 24

os.environ['TF_DETERMINISTIC_OPS'] = 'true'
os.environ['PYTHONHASHSEED'] = f'{seed}'
np.random.seed(seed)
random.seed(seed)
tf.random.set_seed(seed)
sample(seed)

os.environ['TF_DETERMINISTIC_OPS'] = 'true'
os.environ['PYTHONHASHSEED'] = f'{seed}'
np.random.seed(seed)
random.seed(seed)
tf.random.set_seed(seed)
sample(seed)

[csuter]

Great question! TL;DR: pass a Tensor-valued seed (eg tf.constant([24, 42])) or a list of 2 ints (eg, [24, 42]) as the seed to sample_chain to activate "stateless" samplers under the hood.

Longer answer:
As of recently it is possible, but documentation is probably trailing here. Old-style samplers in TF were "stateful" and behaved differently, in subtle ways, between graph mode, eager mode, and XLA mode. Reproducible sampling was possible (but tricky) in graph mode and eager mode, but completely impossible in XLA mode. There are now TF samplers that are stateless but require more careful seed management. We still support the old way, for backward compatibility, which you will get if you pass None or an int seed. The stateless sampling is enabled by passing an int-pair as seed, either a python list of 2 ints or a tf.Tensor of dtype tf.int32 and shape [2]. If you just do this, you should be able to omit all the other global seed-setting lines in your code example. I think nearly all TFP Distributions have been updated to support stateless sampling. JointDistribution variants were updated pretty recently, so will only work if you're using the tfp-nightly pip package. We will have these out in the next stable release as well.

[csuter]

Just to clarify, setting experimental_compile=True is what triggers XLA mode, wherein the old stateful samplers cannot be made to behave reproducibly (but new stateless samplers can).

[brianwa84]

+1 to Chris's detailed answer. Some examples in https://www.tensorflow.org/probability/examples/TFP_Release_Notebook_0_11_0#similarly_tfpmcmcsample_chain_now_accepts_stateless_seeds