Inconsistent usage of STS Parameter bijector

Documentation in structural_time_series.py indicates that the bijector in the Parameter named tuple type is intended to transform unconstrained values into the support space of the parameter.

However, in Autoregressive, the bijector is used with the default MultivariateNormalDiag in a different (and seemingly incorrect) way.

Autoregressive allows specification of coefficient_constraining_bijector which is intended to constrain the values of the coefficient parameter to a specified subspace.

Unfortunately, this is inconsistent with the usage of Parameter in fitting.py and structural_time_series.py, since it seems to be implicitly assumed that the image of the bijector and domain of prior are equivalent.

Note the lack of use of bijector in StructuralTimeSeries::prior_sample and StructuralTimeSeries::joint_log_prob.

In fitting.py, however, the bijector is applied both in sample_uniform_initial_state and fit_with_hmc.

This gives rise to a problem in the case of Autoregressive (with default prior/constraint) because the support of the prior and the inverse-bijector do not match. This produces NaN/inf values as a result.

This seems to be a design issue with the Parameter type and Autoregressive. One of them has to change, but I would like the maintainers to weigh in which one is using the Parameter bijector as intended.

Simple example to reproduce this issue:

import tensorflow as tf
import tensorflow_probability as tfp
tf.debugging.enable_check_numerics()
ar10 = tfp.sts.Autoregressive(10)
seasonal = tfp.sts.Seasonal(12)
model = tfp.sts.Sum([ar10, seasonal])
variational_posterior = tfp.sts.build_factored_surrogate_posterior(model)

which results in this error:

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-10-2a332d89e5e3> in <module>
----> 1 variational_posterior = tfp.sts.build_factored_surrogate_posterior(model)

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/sts/fitting.py in build_factored_surrogate_posterior(model, batch_shape, seed, name)
    191     for param in model.parameters:
    192       variational_posterior[param.name] = _build_posterior_for_one_parameter(
--> 193           param, batch_shape=batch_shape, seed=seed())
    194     return joint_distribution_named_lib.JointDistributionNamed(
    195         variational_posterior)

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/sts/fitting.py in _build_posterior_for_one_parameter(param, batch_shape, seed)
     77   initial_loc = sample_uniform_initial_state(
     78       param, init_sample_shape=batch_shape,
---> 79       return_constrained=False, seed=seed)
     80   loc = tf.Variable(initial_value=initial_loc,
     81                     name=param.name + '_loc')

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/sts/fitting.py in sample_uniform_initial_state(parameter, return_constrained, init_sample_shape, seed)
     60   """
     61   unconstrained_prior_sample = parameter.bijector.inverse(
---> 62       parameter.prior.sample(init_sample_shape, seed=seed))
     63   uniform_initializer = 4 * tf.random.uniform(
     64       tf.shape(unconstrained_prior_sample),

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/bijectors/bijector.py in inverse(self, y, name, **kwargs)
   1071       NotImplementedError: if `_inverse` is not implemented.
   1072     """
-> 1073     return self._call_inverse(y, name, **kwargs)
   1074 
   1075   def _compute_inverse_log_det_jacobian_with_caching(

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/bijectors/bijector.py in _call_inverse(self, y, name, **kwargs)
   1043       if mapping.x is not None:
   1044         return mapping.x
-> 1045       mapping = mapping.merge(x=self._inverse(y, **kwargs))
   1046       # It's most important to cache the x->y mapping, because computing
   1047       # forward(inverse(y)) may be numerically unstable / lossy. Caching the

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_probability/python/bijectors/tanh.py in _inverse(self, y)
     59 
     60   def _inverse(self, y):
---> 61     return tf.atanh(y)
     62 
     63   # We implicitly rely on _forward_log_det_jacobian rather than explicitly

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_math_ops.py in atanh(x, name)
   1314       _result = _pywrap_tensorflow.TFE_Py_FastPathExecute(
   1315         _ctx._context_handle, tld.device_name, "Atanh", name,
-> 1316         tld.op_callbacks, x)
   1317       return _result
   1318     except _core._FallbackException:

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/debug/lib/check_numerics_callback.py in callback(***failed resolving arguments***)
    282                   slot, len(outputs), op_type, output, inputs,
    283                   stack_height_limit=self._stack_height_limit,
--> 284                   path_length_limit=self._path_length_limit))
    285 
    286 

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_array_ops.py in check_numerics(tensor, message, name)
    900         raise
    901     except _core._NotOkStatusException as e:
--> 902       _ops.raise_from_not_ok_status(e, name)
    903   # Add nodes to the TensorFlow graph.
    904   message = _execute.make_str(message, "message")

~/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py in raise_from_not_ok_status(e, name)
   6604   message = e.message + (" name: " + name if name is not None else "")
   6605   # pylint: disable=protected-access
-> 6606   six.raise_from(core._status_to_exception(e.code, message), None)
   6607   # pylint: enable=protected-access
   6608 

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

InvalidArgumentError: 

!!! Detected Infinity or NaN in output 0 of eagerly-executing op "Atanh" (# of outputs: 1) !!!
  dtype: <dtype: 'float32'>
  shape: (10,)
  # of +NaN elements: 3

  Input tensor: tf.Tensor(
[-1.3287405  -0.18918735  0.9701094  -1.012463    0.5736228   0.57741284
  0.739421   -0.96497154 -1.3628513   0.5749111 ], shape=(10,), dtype=float32)

 : Tensor had NaN values [Op:CheckNumerics

This appears to cause issues during optimization without check_numerics enabled.

Start by reading the Parameter handling in structural_time_series.py, especially prior_sample and joint_log_prob, then trace bijector use in fitting.py, including sample_uniform_initial_state and fit_with_hmc. Reproduce the failure with Autoregressive, Seasonal, and build_factored_surrogate_posterior. Done means the bijector semantics are made consistent and the example no longer produces NaN or infinity values.