Mixture returning NaN `log_prob` for valid sample (or: Softplus(Normal) returning `nan` probability on negative numbers)

I am working on some code for algorithmic trading. I'd like to create a probabilistic model over "take profit & stop loss" combinations. Those are valid if either:

• a) take profit is above current close, and stop loss is below current close,
• b) take profit is below current close, and stop loss is above current close.

(i.e., they cannot be on the same side).

This is a simplified version of my model, which makes a mixture distribution centered on (0, 0):

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
dtype = tf.float64

long_d = tfd.MultivariateNormalDiag(loc=tf.zeros((1, 2,), dtype=dtype))
long_tp_sl = tfd.TransformedDistribution(
    long_d,
    tfb.Chain([
        tfb.Blockwise([
            tfb.Identity(),
            tfb.Scale(tf.constant(-1.0, dtype=dtype)),
        ],
                      block_sizes=[1, 1]),
        tfb.Softplus(),
    ]))

short_d = tfd.MultivariateNormalDiag(loc=tf.zeros((1, 2,), dtype=dtype))
short_tp_sl = tfd.TransformedDistribution(
    short_d,
    tfb.Chain([
        tfb.Blockwise([
            tfb.Scale(tf.constant(-1.0, dtype=dtype)),
            tfb.Identity(),
        ],
                      block_sizes=[1, 1]),
        tfb.Softplus(),
    ]))

trained_tp_sl = tfd.Mixture(
    cat=tfd.Categorical(probs=tf.constant([[0.5, 0.5]], dtype=dtype)),
    components=[long_tp_sl, short_tp_sl])

sample = trained_tp_sl.sample(500)
take_profit = sample[:, 0, 0]
stop_loss = sample[:, 0, 1]

plt.figure()
plt.scatter(x=take_profit, y=stop_loss, s=2)

This produces correct samples, as expected:

However, the log_prob of the returned sample is full of nans:

>>> trained_tp_sl.log_prob(sample)
<tf.Tensor: shape=(500, 1), dtype=float64, numpy=
array([[nan],
       [nan],
       [nan],
...

I tried to reproduce this, and when I just take a mixture of 2 disjoint Uniform distributions, prob and log_prob seem to work fine:

import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_probability as tfp

tfd = tfp.distributions

a = tfd.Uniform(low=-2.0, high=-1.0)
b = tfd.Uniform(low=1.0, high=2.0)
mix = tfd.Mixture(
    cat=tfd.Categorical(probs=[0.5, 0.5]),
    components=[a, b])

x = tf.linspace(-3.0, 3.0, num=100)

plt.figure()
plt.plot(x, a.prob(x), x, b.prob(x))

plt.figure()
plt.plot(x, mix.prob(x))

print(a.prob(1.5))  # -> 0.0
print(b.prob(1.5))  # -> 1.0
print(mix.prob(1.5))  # -> 0.5, OK
print()
print(a.log_prob(1.5))  # -> -inf
print(b.log_prob(1.5))  # -> 0
print(mix.log_prob(1.5))  # -> -0.6931472, OK

But I noticed that one my disjoint distributions return nan on samples where their probability is 0, instead of -inf:

>>> long_tp_sl.log_prob(sample)
<tf.Tensor: shape=(500, 1), dtype=float64, numpy=
array([[        nan],  # sample from short_tp_sl
       [        nan],
       [        nan],
       [        nan],
       [        nan],
       [-0.28825895],
       [-2.07584651],
       [        nan],
...

(On the other hand, when I tried this on the mix of disjoint Uniform distributions, notice the log_prob on samples with 0 probability was -inf.)

In fact, Normal distribution transformed with a Softplus returns nan probability on negative numbers:

>>> d = tfd.TransformedDistribution(tfd.Normal(loc=0.0, scale=1.0), tfb.Softplus())
>>> print(d.prob(-1.0))
tf.Tensor(nan, shape=(), dtype=float32)

The same also occurs with the Exp bijector.

Is there any way I can fix up my code to make the mixture's log_prob (and prob) behave correctly?

Reproduce the NaN behavior with the provided Softplus, Exp, TransformedDistribution, and Mixture examples. Start by tracing how transformed distributions evaluate log_prob outside the bijector's support, then compare the result with the disjoint Uniform example. Done means invalid component evaluations contribute -inf rather than NaN, while valid samples retain correct probability and log_prob values.