Bijector log_prob outputs nan when executing with @tf.function but not without it.

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS Catalina 10.15.2 (19C57)
• TensorFlow installed from (source or binary): binary
• TensorFlow versions:

$ python -c "import tensorflow as tf; import tensorflow_probability as tfp; print(tf.version.GIT_VERSION, tf.version.VERSION, tfp.__version__)"
v1.12.1-26458-gc251e83805 2.2.0-rc0 0.10.0-dev20200314

Issue
We have a case where we use tfp bijectors to transform a latent gaussian distribution into another conditioned on inputs (reinforcement learning observations specifically). The shift and scale parameters of the transformation are parameterized by a feedforward network that takes the observations as inputs. The output of the transformed gaussian distribution is finally passed through a Tanh bijector.

As shown below, the implementation uses two custom bijectors (ConditionalScale and ConditionalShift) to handle the transformation. The reason for this is that I've been unable to implement the same functionality cleanly with the existing Scale and Shift bijectors [1, 2].

Now, this implementation works perfectly fine when running things in without tf.function decorators, or alternatively, as demonstrated below, when running with tf.config.experimental_run_functions_eagerly(True). I don't understand well enough how tf.function does the tracing to be able exactly point out what the root cause of the error below is, but ultimately the issue is that with tf.functions, the code below consistently produces nans, whereas when disabling them things work fine. Indeed, in my original implementation, which can be found in the softlearning repo, the same setup learns fine even when running in graph mode as long as I remove the tf.function decorators from the GaussianPolicy.{actions,log_probs} methods.

My questions thus are:

• Am I doing something wrong with the tf.function decorators? Currently, it seems like a bug to be, but I'm unsure about that because of my lack of experience with how the tfp objects work together with tf.functions.
• If it's not a bug, does anyone have thoughts about what the best practice of implementing such a model is? Currently, it seems logical to wrap the operations (i.e. shift/scale computation and the subsequent sampling based on them) in GaussianPolicy.{actions,log_probs} inside a tf.function, but maybe that's wrong and there's actually a better way to achieve similar effect? Previously, our GaussianPolicy was implemented fully with tf.keras.Models which turned out to be super messy.
• If it's not a bug, then it would obviously be nice to get the issue fixed, but I'd also like to understand what the root cause for it is. The only guess I have right now is that somehow self.shift_and_scale_model or self.action_distribution in GaussianPolicy.{actions,log_probs} get traced incorrectly withing the tf.function.

Standalone code to reproduce the issue
Sorry, the code is a bit verbose because I didn't really understand where the issue arises exactly and thus couldn't prune it down much more. This is a slightly stripped version of my implementation in the softlearning project.

import sys

import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.bijectors import bijector
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import dtype_util


class ConditionalScale(bijector.Bijector):
    def __init__(self,
                 dtype=tf.float32,
                 validate_args=False,
                 name='conditional_scale'):
        """Instantiates the `ConditionalScale` bijector.

        This `Bijector`'s forward operation is:

        ```none
        Y = g(X) = scale * X
        ```

        Args:
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.
        """
        parameters = dict(locals())
        with tf.name_scope(name) as name:
            super(ConditionalScale, self).__init__(
                forward_min_event_ndims=0,
                is_constant_jacobian=True,
                validate_args=validate_args,
                dtype=dtype,
                parameters=parameters,
                name=name)

    def _maybe_assert_valid_scale(self, scale):
        if not self.validate_args:
            return ()
        is_non_zero = assert_util.assert_none_equal(
            scale,
            tf.zeros((), dtype=scale.dtype),
            message='Argument `scale` must be non-zero.')
        return (is_non_zero, )

    def _forward(self, x, scale):
        with tf.control_dependencies(self._maybe_assert_valid_scale(scale)):
            return x * scale

    def _inverse(self, y, scale):
        with tf.control_dependencies(self._maybe_assert_valid_scale(scale)):
            return y / scale

    def _forward_log_det_jacobian(self, x, scale):
        with tf.control_dependencies(self._maybe_assert_valid_scale(scale)):
            return tf.math.log(tf.abs(scale))


class ConditionalShift(bijector.Bijector):
    """Compute `Y = g(X; shift) = X + shift`.

    where `shift` is a numeric `Tensor`.

    Example Use:

    ```python
    shift = Shift([-1., 0., 1])
    x = [1., 2, 3]
    # `forward` is equivalent to:
    # y = x + shift
    y = shift.forward(x)  # [0., 2., 4.]
    ```

    """
    def __init__(self,
                 dtype=tf.float32,
                 validate_args=False,
                 name='conditional_shift'):
        """Instantiates the `ConditionalShift` bijector.

        Args:
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.
        """
        parameters = dict(locals())
        with tf.name_scope(name) as name:
            super(ConditionalShift, self).__init__(
                forward_min_event_ndims=0,
                is_constant_jacobian=True,
                dtype=dtype,
                validate_args=validate_args,
                parameters=parameters,
                name=name)

    @classmethod
    def _is_increasing(cls):
        return True

    def _forward(self, x, shift):
        return x + shift

    def _inverse(self, y, shift):
        return y - shift

    def _forward_log_det_jacobian(self, x, shift):
        # is_constant_jacobian = True for this bijector, hence the
        # `log_det_jacobian` need only be specified for a single input, as this will
        # be tiled to match `event_ndims`.
        return tf.zeros((), dtype=dtype_util.base_dtype(x.dtype))


class GaussianPolicy():
    def __init__(self, input_shape, output_shape):
        output_size = tf.reduce_prod(output_shape)

        input_ = tf.keras.layers.Input(input_shape)
        out = input_
        out = tf.keras.layers.Dense(10, activation='relu')(out)
        out = tf.keras.layers.Dense(10, activation='relu')(out)
        out = tf.keras.layers.Dense(output_size * 2, activation='linear')(out)

        def split_shift_and_log_scale_diag_fn(inputs):
            shift_and_log_scale_diag = inputs
            shift, log_scale_diag = tf.split(
                shift_and_log_scale_diag,
                num_or_size_splits=2,
                axis=-1)
            scale_diag = tf.exp(log_scale_diag)
            return [shift, scale_diag]

        shift, scale = tf.keras.layers.Lambda(
            split_shift_and_log_scale_diag_fn
        )(out)

        self.shift_and_scale_model = tf.keras.Model(input_, (shift, scale))

        base_distribution = tfp.distributions.MultivariateNormalDiag(
            loc=tf.zeros(output_shape), scale_diag=tf.ones(output_shape))

        raw_action_distribution = tfp.bijectors.Chain((
            ConditionalShift(name='shift'),
            ConditionalScale(name='scale'),
        ))(base_distribution)

        self.base_distribution = base_distribution
        self.raw_action_distribution = raw_action_distribution
        self.action_distribution = tfp.bijectors.Tanh()(
            raw_action_distribution)

    @tf.function(experimental_relax_shapes=True)
    def actions(self, observations):
        """Compute actions for given observations."""

        batch_shape = tf.shape(observations)[0]
        shifts, scales = self.shift_and_scale_model(observations)
        actions = self.action_distribution.sample(
            batch_shape,
            bijector_kwargs={'scale': {'scale': scales},
                             'shift': {'shift': shifts}})

        return actions

    @tf.function(experimental_relax_shapes=True)
    def log_probs(self, observations, actions):
        """Compute log probabilities of `actions` given observations."""

        shifts, scales = self.shift_and_scale_model(observations)
        log_probs = self.action_distribution.log_prob(
            actions,
            bijector_kwargs={'scale': {'scale': scales},
                             'shift': {'shift': shifts}}
        )[..., tf.newaxis]

        return log_probs


def main():
    # NOTE: this is using a fixed input. However, the issue is consistently reproducible with
    # with "real" observations coming from the RL environment, i.e. the eager version
    # of the code never fails, whereas the graph version fails randomly (yet consistently).
    observations = tf.repeat([[
        0.23420376, -0.32872833, 0.03206815, 0.18556681,
        -2.0855187, 2.12688574, -3.91442398, 2.80974896,
    ]], 256, axis=0)  # (256, 8)

    found_nans = False

    for i in range(10):
        # Loop over a few different initialization of policy parameters.
        policy = GaussianPolicy(
            input_shape=observations.shape[1:], output_shape=(2, ))

        actions = policy.actions(observations)  # (256, 2)
        log_probs = policy.log_probs(observations, actions)  # (256, 1)

        if tf.reduce_any(tf.math.is_nan(log_probs)):
            found_nans = True
            break

    assert not found_nans, "Failure."

    print("Success.")


if __name__ == '__main__':
    run_eagerly = sys.argv[1].lower() == 'true'
    tf.config.experimental_run_functions_eagerly(run_eagerly)
    main()

Other info / logs

$ python -m tests.test_broken_actions_v2 True
Success.

$ python -m tests.test_broken_actions_v2 False
Traceback (most recent call last):
  File "/Users/hartikainen/conda/envs/softlearning-2/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/Users/hartikainen/conda/envs/softlearning-2/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/Users/hartikainen/github/rail-berkeley/softlearning-2/tests/test_broken_actions_v2.py", line 213, in <module>
    main()
  File "/Users/hartikainen/github/rail-berkeley/softlearning-2/tests/test_broken_actions_v2.py", line 205, in main
    assert not found_nans, "Failure."
AssertionError: Failure.

Start by running tests/test_broken_actions_v2.py with eager execution enabled and disabled, then inspect GaussianPolicy.actions and log_probs together with the ConditionalScale and ConditionalShift bijectors. Compare the traced and eager paths around action_distribution.sample and log_prob; done means the NaN cause is identified and the graph-mode behavior is corrected or clearly documented with a reliable reproduction.