tensorflow / probability

Probabilistic reasoning and statistical analysis in TensorFlow
https://www.tensorflow.org/probability/
Apache License 2.0
4.24k stars 1.09k forks source link

Unable to sample a multilevel model #601

Closed fernandoramacciotti closed 4 years ago

fernandoramacciotti commented 4 years ago

Hi everyone,

I am trying to adapt a PyMC3 multilevel model to TFP for predicting soccer matches. It is a partial pooled model, similar to the radon example, so it basically estimate scoring distribution via Poisson for each of the 20 teams, from 203 datapoints (matches).

I wasn't sure how to build the whole model in a modular fashion, so I included everything into a single tfp.JointDistributionSequential call.

def home_goals():
    return tfd.JointDistributionSequential([
        # home attack rate
        tfd.Normal(loc=0., scale=1.),           # mu                       
        tfd.Gamma(concentration=2., rate=2.),   # sd
        tfd.Independent(                        # nc
            tfd.Normal(
                loc=tf.zeros([num_home_teams]),
                scale=tf.ones([num_home_teams])),
            reinterpreted_batch_ndims=None),
        lambda anc, asd, amu: tfd.Independent(     # attack rate
            tfd.Deterministic(
                amu[..., tf.newaxis] + anc * asd[..., tf.newaxis]),
            reinterpreted_batch_ndims=None),
        # away defense rate
        tfd.Gamma(concentration=2., rate=2.),   # sd
        tfd.Independent(                        # nc
            tfd.Normal(
                loc=tf.zeros([num_away_teams]),
                scale=tf.ones([num_away_teams])),
            reinterpreted_batch_ndims=1),
        lambda dnc, dsd: tfd.Independent(         # defense rate
            tfd.Deterministic(
                dnc * dsd[..., tf.newaxis]),
            reinterpreted_batch_ndims=1),
        # home attack advantage
        tfd.Normal(loc=0., scale=1.),           # mu                       
        tfd.Gamma(concentration=2., rate=2.),   # sd
        tfd.Independent(                        # nc
            tfd.Normal(
                loc=tf.zeros([num_home_teams]),
                scale=tf.ones([num_home_teams])),
            reinterpreted_batch_ndims=1),
        lambda hanc, hasd, hamu: tfd.Independent(     # home rate
            tfd.Deterministic(
                hamu[..., tf.newaxis] + hanc * hasd[..., tf.newaxis]),
            reinterpreted_batch_ndims=1),
        # home diff
        lambda har, hnc, hsd, hmu, dr, dnc, dsd, ar: tfd.Independent(
            tfd.Deterministic(ar - dr + har),
            reinterpreted_batch_ndims=1),
        # home goals
        lambda homediff: tfd.Independent(
            tfd.Poisson(rate=tf.exp(tf.gather(homediff, home_team, axis=-1))),
            reinterpreted_batch_ndims=1),  
    ])

def logp(amu, asd, anc, dsd, dnc, hamu, hasd, hanc):
    return home_goals().log_prob([
        amu, asd, anc, dsd, dnc, hamu, hasd, hanc, home_score])

My settings for mcmc sampling are as bellow:

# test number 1 (same error)
initial_state = [
    # home diff
    tf.zeros([num_chains]),# name='init_att_mu_home'),
    tf.ones([num_chains]),# name='init_att_sd_home'),
    tf.zeros([num_chains, num_teams]),# name='init_att_nc_home'),
    tf.zeros([num_chains]),# name='init_def_sd_away'),
    tf.ones([num_chains, num_teams]),# name='init_def_nc_away'),
    tf.zeros([num_chains]),# name='init_homeatt_mu'),
    tf.ones([num_chains]), # name='init_homeatt_sd'),
    tf.ones([num_chains, num_teams]),# name='init_homeatt_nc'),
]

# test 2 (same error)
#initial_state = [home_goals().sample(num_chains)[i] for i in [0, 1, 2, 4, 5, 7, 8, 9]]

# constraints
unconstraining_bijectors = [
    # home diff
    tfb.Identity(),
    tfb.Exp(),
    tfb.Identity(),
    tfb.Identity(),
    tfb.Exp(),
    tfb.Identity(),
    tfb.Exp(),
    tfb.Identity(),
]

kernel = tfp.mcmc.TransformedTransitionKernel(
    inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=logp,
        num_leapfrog_steps=10,
        step_size=0.01),
    bijector=unconstraining_bijectors)

samples, _ = tfp.mcmc.sample_chain(
    num_results=2000, 
    num_burnin_steps=500,
    kernel=kernel,
    current_state=initial_state,
    parallel_iterations=1)

But it throws an incompatible-shape-error (full traceback bellow). I cannot see where the [203,20] tensor is generated.

Can anyone assist me on this, please?

Thanks in advance, Fernando

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-181-25cd018fdb32> in <module>
     19     kernel=kernel,
     20     current_state=initial_state,
---> 21     parallel_iterations=1)

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/sample.py in sample_chain(num_results, current_state, previous_kernel_results, kernel, num_burnin_steps, num_steps_between_results, trace_fn, return_final_kernel_results, parallel_iterations, name)
    324         current_state)
    325     if previous_kernel_results is None:
--> 326       previous_kernel_results = kernel.bootstrap_results(current_state)
    327 
    328     if trace_fn is None:

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/transformed_kernel.py in bootstrap_results(self, init_state, transformed_init_state)
    344           transformed_state=transformed_init_state,
    345           inner_results=self._inner_kernel.bootstrap_results(
--> 346               transformed_init_state))
    347       return kernel_results

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/hmc.py in bootstrap_results(self, init_state)
    558   def bootstrap_results(self, init_state):
    559     """Creates initial `previous_kernel_results` using a supplied `state`."""
--> 560     kernel_results = self._impl.bootstrap_results(init_state)
    561     if self.step_size_update_fn is not None:
    562       step_size_assign = self.step_size_update_fn(self.step_size, None)  # pylint: disable=not-callable

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/metropolis_hastings.py in bootstrap_results(self, init_state)
    262         name=mcmc_util.make_name(self.name, 'mh', 'bootstrap_results'),
    263         values=[init_state]):
--> 264       pkr = self.inner_kernel.bootstrap_results(init_state)
    265       if not has_target_log_prob(pkr):
    266         raise ValueError(

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/hmc.py in bootstrap_results(self, init_state)
    770           init_target_log_prob,
    771           init_grads_target_log_prob,
--> 772       ] = mcmc_util.maybe_call_fn_and_grads(self.target_log_prob_fn, init_state)
    773       if self._store_parameters_in_results:
    774         return UncalibratedHamiltonianMonteCarloKernelResults(

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/internal/util.py in maybe_call_fn_and_grads(fn, fn_arg_list, result, grads, check_non_none_grads, name)
    231     fn_arg_list = (list(fn_arg_list) if is_list_like(fn_arg_list)
    232                    else [fn_arg_list])
--> 233     result, grads = _value_and_gradients(fn, fn_arg_list, result, grads)
    234     if not all(r.dtype.is_floating
    235                for r in (result if is_list_like(result) else [result])):  # pylint: disable=superfluous-parens

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/internal/util.py in _value_and_gradients(fn, fn_arg_list, result, grads, name)
    190 
    191     if result is None:
--> 192       result = fn(*fn_arg_list)
    193       if grads is None and tf.executing_eagerly():
    194         # Ensure we disable bijector cacheing in eager mode.

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/transformed_kernel.py in new_target_log_prob(*transformed_state_parts)
    195       ]
    196       tlp = target_log_prob_fn(
--> 197           *self._forward_transform(transformed_state_parts))
    198       event_ndims = [
    199           tf.rank(sp) - tf.rank(tlp) for sp in transformed_state_parts

<ipython-input-147-a36e1604395c> in logp(amu, asd, anc, dsd, dnc, hamu, hasd, hanc)
     30 
     31 def logp(amu, asd, anc, dsd, dnc, hamu, hasd, hanc):
---> 32     return home_goals().log_prob([amu, asd, anc, dsd, dnc, hamu, hasd, hanc, home_score])

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/distributions/distribution.py in log_prob(self, value, name, **kwargs)
    872         values of type `self.dtype`.
    873     """
--> 874     return self._call_log_prob(value, name, **kwargs)
    875 
    876   def _call_prob(self, value, name, **kwargs):

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/distributions/joint_distribution.py in _call_log_prob(self, value, name)
    320   def _call_log_prob(self, value, name):
    321     with self._name_and_control_scope(name):
--> 322       return self._log_prob(value)
    323 
    324   def _call_sample_n(self, sample_shape, seed, name, value=None):

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/distributions/joint_distribution.py in _log_prob(self, value)
    275 
    276   def _log_prob(self, value):
--> 277     _, xs = self._map_measure_over_dists('log_prob', value)
    278     return sum(xs)
    279 

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/distributions/joint_distribution.py in _map_measure_over_dists(self, attr, value)
    291     if any(x is None for x in tf.nest.flatten(value)):
    292       raise ValueError('No `value` part can be `None`; saw: {}.'.format(value))
--> 293     ds, xs = self._call_flat_sample_distributions(value=value)
    294     return ds, maybe_check_wont_broadcast(
    295         (getattr(d, attr)(x) for d, x in zip(ds, xs)),

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/distributions/joint_distribution.py in _call_flat_sample_distributions(self, sample_shape, seed, value)
    309     if value is not None:
    310       value = self._model_flatten(value)
--> 311     ds, xs = self._flat_sample_distributions(sample_shape, seed, value)
    312     if not value:
    313       self._most_recently_built_distributions = tuple(ds)

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/distributions/joint_distribution_sequential.py in _flat_sample_distributions(self, sample_shape, seed, value)
    234     for i, (dist_fn, args) in enumerate(zip(self._dist_fn_wrapped,
    235                                             self._dist_fn_args)):
--> 236       ds.append(dist_fn(*xs[:i]))  # Chain rule of probability.
    237       if xs[i] is None:
    238         # TODO(b/129364796): We should ignore args prefixed with `_`; this

~/anaconda3/lib/python3.7/site-packages/tensorflow_probability/python/distributions/joint_distribution_sequential.py in dist_fn_wrapped(*xs)
    449           '(dist_fn: {}, expected: {}, saw: {}).'.format(
    450               i, dist_fn, len(args), len(xs)))
--> 451     return dist_fn(*reversed(xs[-len(args):]))
    452   return dist_fn_wrapped, args
    453 

<ipython-input-174-c2a302fcabf2> in <lambda>(hanc, hasd, hamu)
     34         lambda hanc, hasd, hamu: tfd.Independent(     # home rate
     35             tfd.Deterministic(
---> 36                 hamu[..., tf.newaxis] + hanc * hasd[..., tf.newaxis]),
     37             reinterpreted_batch_ndims=1),
     38         # home diff

~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/math_ops.py in binary_op_wrapper(x, y)
    897     with ops.name_scope(None, op_name, [x, y]) as name:
    898       if isinstance(x, ops.Tensor) and isinstance(y, ops.Tensor):
--> 899         return func(x, y, name=name)
    900       elif not isinstance(y, sparse_tensor.SparseTensor):
    901         try:

~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/math_ops.py in _add_dispatch(x, y, name)
   1195       return gen_math_ops.add(x, y, name=name)
   1196     else:
-> 1197       return gen_math_ops.add_v2(x, y, name=name)
   1198   else:
   1199     return gen_math_ops.add(x, y, name=name)

~/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_math_ops.py in add_v2(x, y, name)
    544       else:
    545         message = e.message
--> 546       _six.raise_from(_core._status_to_exception(e.code, message), None)
    547   # Add nodes to the TensorFlow graph.
    548   _, _, _op = _op_def_lib._apply_op_helper(

~/anaconda3/lib/python3.7/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Incompatible shapes: [1,20,1] vs. [203,20] [Op:AddV2] name: mcmc_sample_chain/transformed_kernel_bootstrap_results/mh_bootstrap_results/hmc_kernel_bootstrap_results/maybe_call_fn_and_grads/value_and_gradients/mcmc_sample_chain_transformed_kernel_bootstrap_results_mh_bootstrap_results_hmc_kernel_bootstrap_results_maybe_call_fn_and_grads_value_and_gradients_JointDistributionSequential/log_prob/add/
junpenglao commented 4 years ago

The error you are seeing is because the shape in initial_state is incorrect. It originated from the used of tfd.Deterministic. Since it is a distribution it will become an input to the logp which means you also need to supply the value in the initial_state, however it will also not work as there will be no gradient in HMC. Overall, tfd.Deterministic is not the same as pm.Deterministic, which means you cannot use it to wrap a Deterministic transformation.

If you could share a notebook with data, I can take a look for you help cleaning it up the implementation a bit :-)

fernandoramacciotti commented 4 years ago

Hi @junpenglao, I really appreciate your help. Could you explore further the differences between tfd.Deterministic and pm.Deterministic? Also, in the article I shared, both home_goals and away_goals are declared inside the same pm.model object and, therefore, estimated together. Is it possible to do it in TFP?

Please find the notebook URL. Thanks again =)

junpenglao commented 4 years ago

Could you explore further the differences between tfd.Deterministic and pm.Deterministic?

pm.Deterministic is a way to keep track of deterministic transformation in a pm.Model (meaning that pymc3 will actively log these values in the sample trace). It does not modify the log_prob and it does not add additional free parameters to the model. Instead, tfd.Deterministic is a distribution with dirac measure, which means that it actually adds a new free parameters to the model. I am not completely sure about the intended use case for it, maybe @jvdillon could comment a bit more.

Also, in the article I shared, both home_goals and away_goals are declared inside the same pm.model object and, therefore, estimated together. Is it possible to do it in TFP?

Yes it is possible, you just need to add one more distribution.

I will have a look at the notebook and get back to you.

junpenglao commented 4 years ago

Since there are quite some transformation in the model, using JointDistributionCoroutine might be a bit easier:

Root = tfd.JointDistributionCoroutine.Root
def model():  # <== need to be a model with no input and no return
  # Home attack rate
  attack_hyper = yield Root(tfd.Sample(tfd.Normal(loc=0., scale=1.), 1))
  attack_hyper_sd = yield Root(tfd.Sample(tfd.Gamma(concentration=2., rate=2.), 1))
  attack_rate_nc = yield Root(tfd.Sample(tfd.Normal(loc=0., scale=1.), num_home_teams))
  attack_rate = attack_hyper + attack_rate_nc * attack_hyper_sd
  # Away defense rate
  defense_hyper_sd = yield Root(tfd.Sample(tfd.Gamma(concentration=2., rate=2.), 1))
  defense_rate_nc = yield Root(tfd.Sample(tfd.Normal(loc=0., scale=1.), num_away_teams))
  defense_rate = defense_rate_nc * defense_hyper_sd
  # Home attack advantage
  home_attack_hyper_sd = yield Root(tfd.Sample(tfd.Gamma(concentration=2., rate=2.), 1))
  home_attack_hyper = yield Root(tfd.Sample(tfd.Normal(loc=0., scale=1.), 1))
  home_attack_nc = yield Root(tfd.Sample(tfd.Normal(loc=0., scale=1.), num_home_teams))
  home_attack_advantage = home_attack_hyper + home_attack_nc * home_attack_hyper_sd
  # Home defense advantage
  home_defense_hyper_sd = yield Root(tfd.Sample(tfd.Gamma(concentration=2., rate=2.), 1))
  home_defense_hyper = yield Root(tfd.Sample(tfd.Normal(loc=0., scale=1.), 1))
  home_defense_nc = yield Root(tfd.Sample(tfd.Normal(loc=0., scale=1.), num_home_teams))
  home_defense_advantage = home_defense_hyper + home_defense_nc * home_defense_hyper_sd
  # Liklihood
  home_diff = tf.gather(attack_rate, home_team, axis=-1) - \
              tf.gather(defense_rate, away_team, axis=-1) + \
              tf.gather(home_attack_advantage, home_team, axis=-1)

  away_diff = tf.gather(attack_rate, away_team, axis=-1) - \
              tf.gather(defense_rate, home_team, axis=-1) - \
              tf.gather(home_defense_advantage, home_team, axis=-1)
  home_goals = yield tfd.Independent(tfd.Poisson(log_rate=home_diff), 1)
  away_goals = yield tfd.Independent(tfd.Poisson(log_rate=away_diff), 1)

model_jd = tfd.JointDistributionCoroutine(model)

You can then check the model log_prob:

test_samples = model_jd.sample(5)
model_jd.log_prob(test_samples)

and create a log_prob function closure (conditioned on the observed) for sampling:

unnomarlized_log_prob = lambda *args: model_jd.log_prob(list(args) + [
  home_score[tf.newaxis, ...], away_score[tf.newaxis, ...]])

num_chains = 5

initial_state = model_jd.sample(num_chains)[:-2]
unnomarlized_log_prob(*initial_state)

You can take a look at my edit of your colab. I am also using the NUTS sampler there :-)

Let me know if there is anything not clear!

fernandoramacciotti commented 4 years ago

thank you very much, @junpenglao , really appreciate it. Very nice clean code! Could you please just explain me when to use tfd.Sample and when not to?

Thanks!

junpenglao commented 4 years ago

Oh right, thanks for pointing that out! So this is more my personal preference, as you see in your own code, you can are doing some_scaler_distribution and later on using rv_from_some_scaler_distribution[..., tf.newaxis]. Using tfd.Sample(some_scaler_distribution, 1) is basically preemptive doing that and make sure boardcast works directly downstream. So the options are:

a = tfd.Sample(tfd.Normal(0., 1.), 1)
a.sample(5)              # <== shape=(5, 1)
a.log_prob(a.sample(5))  # <== shape=(5,)

a = tfd.Independent(tfd.Normal([0.], 1.), 1)
a.sample(5)              # <== shape=(5, 1)
a.log_prob(a.sample(5))  # <== shape=(5,)

a = tfd.Normal(0., 1.)
a.sample(5)              # <== shape=(5,) but need to do rv[..., tf.newaxis]
a.log_prob(a.sample(5))  # <== shape=(5,)

And personally I prefer the first one.

fernandoramacciotti commented 4 years ago

perfect, thank you very much!

Em sex, 11 de out de 2019 18:01, Junpeng Lao notifications@github.com escreveu:

Oh right, thanks for pointing that out! So this is more my personal preference, as you see in your own code, you can are doing some_scaler_distribution and later on using rv_from_some_scaler_distribution[..., tf.newaxis]. Using tfd.Sample(some_scaler_distribution, 1) is basically preemptive doing that and make sure boardcast works directly downstream. So the options are:

a = tfd.Sample(tfd.Normal(0., 1.), 1) a.sample(5) # <== shape=(5, 1) a.log_prob(a.sample(5)) # <== shape=(5,)

a = tfd.Independent(tfd.Normal([0.], 1.), 1) a.sample(5) # <== shape=(5, 1) a.log_prob(a.sample(5)) # <== shape=(5,)

a = tfd.Normal(0., 1.) a.sample(5) # <== shape=(5,) but need to do rv[..., tf.newaxis] a.log_prob(a.sample(5)) # <== shape=(5,)

— You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub https://github.com/tensorflow/probability/issues/601?email_source=notifications&email_token=ADU3DVG2DEJN3HYXJAT3PL3QODSRRA5CNFSM4I7T4E6KYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEBBGNIQ#issuecomment-541222562, or unsubscribe https://github.com/notifications/unsubscribe-auth/ADU3DVFC7HL6MSUATAXCEXDQODSRRANCNFSM4I7T4E6A .

junpenglao commented 4 years ago

Happy to help! Feel free to follow up in this issue or a new one if you have further question :-)

ChaudhuriC commented 1 year ago

I am facing a similar problem in one of my models: ######################################################### import tensorflow as tf import tensorflow_probability as tfp import collections import numpy as np import pandas as pd

tfd=tfp.distributions tfb = tfp.bijectors

data=pd.read_csv('data_fit_sub_1222.csv') num_stations=len(np.unique(data.COMID)) num_obs=len(data)

discharge=tf.convert_to_tensor(data['DISCHARGE'], dtype=tf.float32) station=tf.convert_to_tensor(data['COMID'], dtype=tf.int32) covar=tf.convert_to_tensor(data[['log_drainage_area']], dtype=tf.float32)

def affine(x, kernel_diag, bias=tf.zeros([])): kernel_diag = tf.ones_like(x) kernel_diag bias = tf.ones_like(x) bias return x * kernel_diag + bias Root = tfd.JointDistributionCoroutine.Root def bhm_model(): alpha1 = yield Root(tfd.Sample(tfd.Normal(loc=3.3, scale=1.2), 1)) alpha2 = yield Root(tfd.Sample(tfd.Uniform(low=0.0, high=2.0), 1)) beta1 = yield Root(tfd.Sample(tfd.Normal(loc=2.3, scale=1.0), 1)) beta2 = yield Root(tfd.Sample(tfd.Uniform(low=0.0, high=2.0), 1)) xi_bar = yield Root(tfd.Sample(tfd.Normal(loc=-0.13, scale=0.08), 1)) tau_mu = yield Root(tfd.Sample(tfd.Gamma(concentration=1.2, rate=0.02), 1)) tau_sigma = yield Root(tfd.Sample(tfd.Gamma(concentration=1.0, rate=0.02), 1)) tau_xi = yield Root(tfd.Sample(tfd.Gamma(concentration=0.08, rate=0.009), 1))

mu_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, alpha2[..., tf.newaxis], alpha1[..., tf.newaxis]), scale_identity_multiplier=tau_mu)
sigma_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, beta2[..., tf.newaxis], beta1[..., tf.newaxis]), scale_identity_multiplier=tau_sigma)
xi_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, xi_bar[..., tf.newaxis], 0), scale_identity_multiplier=tau_xi)

mu_norm_sample = yield mu_norm_dist.sample()
sigma_norm_sample = yield sigma_norm_dist.sample()
xi_norm_sample = yield xi_norm_dist.sample()

mu_gev = yield tf.gather(mu_norm_sample, station)
sigma_gev = yield tf.gather(sigma_norm_sample, station)
xi_gev = yield tf.gather(xi_norm_sample, station)

gev_dist = yield Root(tfd.GeneralizedExtremeValue(loc=mu_gev, scale=sigma_gev, concentration=xi_gev), 1)
y_dist = yield tfd.Independent(gev_dist, reinterpreted_batch_ndims=1)

model_jd = tfd.JointDistributionCoroutine(bhm_model) model_jd.sample(5) ###################################################################################### the error trace is here:

Traceback (most recent call last):

File "C:\Users\ChiranjibChaudhuri\AppData\Local\Temp\ipykernel_30796\3417369031.py", line 28, in model_jd.sample(5)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow_probability\python\distributions\distribution.py", line 1201, in sample return self._call_sample_n(sample_shape, seed, **kwargs)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow_probability\python\distributions\joint_distribution.py", line 954, in _call_sample_n return self._sample_n(

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow_probability\python\internal\distribution_util.py", line 1362, in _fn return fn(*args, **kwargs)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow_probability\python\distributions\joint_distribution.py", line 691, in _sample_n self._get_static_distribution_attributes(seed=seed)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow_probability\python\distributions\joint_distribution.py", line 361, in _get_static_distribution_attributes flat_list_of_static_attributes = callable_util.get_output_spec(

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow_probability\python\internal\callable_util.py", line 63, in get_output_spec return tf.function(fn, autograph=False).get_concrete_function(

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow\python\eager\def_function.py", line 1239, in get_concrete_function concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow\python\eager\def_function.py", line 1219, in _get_concrete_function_garbage_collected self._initialize(args, kwargs, add_initializers_to=initializers)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow\python\eager\def_function.py", line 785, in _initialize self._stateful_fn._get_concrete_function_internal_garbage_collected( # pylint: disable=protected-access

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow\python\eager\function.py", line 2523, in _get_concrete_function_internal_garbage_collected graphfunction, = self._maybe_define_function(args, kwargs)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow\python\eager\function.py", line 2760, in _maybe_define_function graph_function = self._create_graph_function(args, kwargs)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow\python\eager\function.py", line 2670, in _create_graph_function func_graph_module.func_graph_from_py_func(

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow\python\framework\func_graph.py", line 1247, in func_graph_from_py_func func_outputs = python_func(*func_args, **func_kwargs)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow\python\eager\def_function.py", line 677, in wrapped_fn out = weak_wrapped_fn().wrapped(*args, **kwds)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow_probability\python\distributions\joint_distribution.py", line 362, in lambda: self._execute_model( # pylint: disable=g-long-lambda

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow_probability\python\distributions\joint_distribution.py", line 1003, in _execute_model next_value, traced_values = sample_and_trace_fn(

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow_probability\python\distributions\joint_distribution.py", line 156, in trace_static_attributes value = dist.sample(seed=seed)

File "C:\Users\ChiranjibChaudhuri\anaconda3\envs\geosapiens-env\lib\site-packages\tensorflow\python\framework\ops.py", line 446, in getattr self.getattribute(name)

AttributeError: 'Tensor' object has no attribute 'sample'

################################################################# Not sure what I am doing wrong. Pretty new to tf_probability though. TIA.

ColCarroll commented 1 year ago

the formatting is pretty rough, but changing these lines:

mu_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, alpha2[..., tf.newaxis], alpha1[..., tf.newaxis]), scale_identity_multiplier=tau_mu)
sigma_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, beta2[..., tf.newaxis], beta1[..., tf.newaxis]), scale_identity_multiplier=tau_sigma)
xi_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, xi_bar[..., tf.newaxis], 0), scale_identity_multiplier=tau_xi)

mu_norm_sample = yield mu_norm_dist.sample()
sigma_norm_sample = yield sigma_norm_dist.sample()
xi_norm_sample = yield xi_norm_dist.sample()

into

mu_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, alpha2[..., tf.newaxis], alpha1[..., tf.newaxis]), scale_identity_multiplier=tau_mu)
sigma_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, beta2[..., tf.newaxis], beta1[..., tf.newaxis]), scale_identity_multiplier=tau_sigma)
xi_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, xi_bar[..., tf.newaxis], 0), scale_identity_multiplier=tau_xi)

mu_norm_sample = yield mu_norm_dist
sigma_norm_sample = yield sigma_norm_dist
xi_norm_sample = yield xi_norm_dist

or even better

mu_norm_sample = yield tfd.MultivariateNormalDiag(loc=affine(covar, alpha2[..., tf.newaxis], alpha1[..., tf.newaxis]), scale_identity_multiplier=tau_mu)
sigma_norm_sample = yield tfd.MultivariateNormalDiag(loc=affine(covar, beta2[..., tf.newaxis], beta1[..., tf.newaxis]), scale_identity_multiplier=tau_sigma)
xi_norm_sample = yield tfd.MultivariateNormalDiag(loc=affine(covar, xi_bar[..., tf.newaxis], 0), scale_identity_multiplier=tau_xi)

should work

ChaudhuriC commented 1 year ago

Hello Colin, Thank you for replying. I changed the suggested lines. But the error remains the same. I also tried to write a sequential join distribution version of that: ''' def bhm_model(): return tfd.JointDistributionSequential([ tfd.Normal(loc=3.3, scale=1.2), # alpha1 tfd.Uniform(low=0.0, high=2.0), # alpha2 tfd.Normal(loc=2.3, scale=1.0), # beta1 tfd.Uniform(low=0.0, high=2.0), # beta2 tfd.Normal(loc=-0.13, scale=0.08), # xi_bar tfd.Gamma(concentration=1.2, rate=0.02), # tau_mu tfd.Gamma(concentration=1.0, rate=0.02), # tau_sigma tfd.Gamma(concentration=0.08, rate=0.009), # tau_xi lambda alpha1_sample, alpha2_sample, beta1_sample, beta2_sample, xi_bar_sample, tau_mu_sample, tau_sigma_sample, tau_xi_sample: ( tfd.MultivariateNormalDiag(loc=affine(covar, alpha2_sample[..., tf.newaxis], alpha1_sample[..., tf.newaxis]),scale_identity_multiplier=tau_mu_sample), # mu_norm_dist tfd.MultivariateNormalDiag(loc=affine(covar, beta2_sample[..., tf.newaxis], beta1_sample[..., tf.newaxis]),scale_identity_multiplier=tau_sigma_sample), # sigma_norm_dist tfd.MultivariateNormalDiag(loc=affine(covar, xi_bar_sample[..., tf.newaxis], 0),scale_identity_multiplier=tau_xi_sample), # xi_norm_dist ), lambda mu_norm_sample, sigma_norm_sample, xi_norm_sample: ( tf.gather(mu_norm_sample, station), # mu_gev tf.gather(sigma_norm_sample, station), # sigma_gev tf.gather(xi_norm_sample, station), # xi_gev ), lambda mu_gev, sigma_gev, xi_gev: tfd.GeneralizedExtremeValue(loc=mu_gev, scale=sigma_gev, concentration=xi_gev), # gev_dist lambda gev_dist: tfd.Independent(gev_dist,reinterpreted_batch_ndims=1) # y_dist ])

model=bhm_model() model.sample(5) '''

This time the error changed to :

''' AttributeError: 'tuple' object has no attribute 'sample'

''''

I think the problem is coming from the tfd.Independent.

ColCarroll commented 1 year ago

Looks like

Root = tfd.JointDistributionCoroutine.Root
def bhm_model():
  alpha1 = yield Root(tfd.Normal(loc=3.3, scale=1.2))
  alpha2 = yield Root(tfd.Uniform(low=0.0, high=2.0))
  beta1 = yield Root(tfd.Normal(loc=2.3, scale=1.0))
  beta2 = yield Root(tfd.Uniform(low=0.0, high=2.0))
  xi_bar = yield Root(tfd.Normal(loc=-0.13, scale=0.08))
  tau_mu = yield Root(tfd.Gamma(concentration=1.2, rate=0.02))
  tau_sigma = yield Root(tfd.Gamma(concentration=1.0, rate=0.02))
  tau_xi = yield Root(tfd.Gamma(concentration=0.08, rate=0.009))

  mu_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, alpha2[..., tf.newaxis], alpha1[..., tf.newaxis]), scale_diag=(tau_mu[..., None] * tf.ones(num_stations)))
  sigma_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, beta2[..., tf.newaxis], beta1[..., tf.newaxis]), scale_diag=tau_sigma[..., None] * tf.ones(num_stations))
  xi_norm_dist = tfd.MultivariateNormalDiag(loc=affine(covar, xi_bar[..., tf.newaxis], 0), scale_diag=tau_xi[..., None] * tf.ones(num_stations))

  mu_norm_sample = yield mu_norm_dist
  sigma_norm_sample = yield sigma_norm_dist
  xi_norm_sample = yield xi_norm_dist

  mu_gev = tf.gather(mu_norm_sample, station, axis=-1)
  sigma_gev = tf.gather(sigma_norm_sample, station, axis=-1)
  xi_gev = tf.gather(xi_norm_sample, station, axis=-1)

  gev_dist = yield tfd.Independent(tfd.GeneralizedExtremeValue(loc=mu_gev, scale=sigma_gev, concentration=xi_gev), reinterpreted_batch_ndims=1)
  # y_dist = yield tfd.Independent(gev_dist, reinterpreted_batch_ndims=1)

model_jd = tfd.JointDistributionCoroutine(bhm_model)
model_jd.sample(5)
ChaudhuriC commented 1 year ago

Thanks! The error is solved but now the tensor dimension is conflicting and giving error. In fact, I cannot even relate to one of the dimension numbers from the error msg. I am attaching the data and the test version of the script.

BHM.zip