# Buggy code
def log_prob(self, x: Tensor) -> Tensor:
    F = self.F
    alpha, beta = self.alpha, self.beta

    def gamma_log_prob(x, alpha, beta):
        return (
            alpha * F.log(beta) - F.gammaln(alpha) + (alpha - 1) * F.log(x) - beta * x
        )

    """
        The gamma_log_prob(x) above returns NaNs for x<=0. Wherever there are NaN in either of the F.where() conditional
        vectors, then F.where() returns NaN at that entry as well, due to its indicator function multiplication: 
        1*f(x) + np.nan*0 = nan, since np.nan*0 return nan. 
        Therefore replacing gamma_log_prob(x) with gamma_log_prob(abs(x) mitigates nan returns in cases of x<=0 without 
        altering the value in cases of x>0. 
        This is a known issue in pytorch as well https://github.com/pytorch/pytorch/issues/12986.
        """
    return F.where(
        x > 0,
        gamma_log_prob(F.abs(x), alpha, beta),
        -np.inf * F.ones_like(x),
    )


# Traceback (gluonts.core.exception.GluonTSUserError)
gluonts 0.5.1.dev102+g88a9832.d20200924
How close is the simulated data to zero? mixture_samples.abs().min() = 0.00029572556

true mixture.components
gluonts.mx.distribution.gaussian.Gaussian(mu=mxnet.nd.array([0.0], dtype=numpy.float32), sigma=mxnet.nd.array([0.20000000298023224], dtype=numpy.float32))
gluonts.mx.distribution.gamma.Gamma(alpha=mxnet.nd.array([9.0], dtype=numpy.float32), beta=mxnet.nd.array([2.0], dtype=numpy.float32))

MWE: gluon-ts MixtureDistributionOutput() of Gamma-Gaussian is able to learn the component weights and arguments for real-valued data
trainer.learning_rate 0.02
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:05<00:00, 171.16it/s, loss=[1.9887046]]
fit_mixture.components
gluonts.mx.distribution.gaussian.Gaussian(mu=mxnet.nd.array([-0.021975774317979813], dtype=numpy.float32), sigma=mxnet.nd.array([0.1961863487958908], dtype=numpy.float32))
gluonts.mx.distribution.gamma.Gamma(alpha=mxnet.nd.array([3.9501376152038574], dtype=numpy.float32), beta=mxnet.nd.array([0.8736407160758972], dtype=numpy.float32))

gluon-ts DeepAR should be able to learn the component weights and arguments of a Gamma-Gaussian MixtureDistributionOutput() for the same data
trainer.learning_rate 1e-05
learning rate from ``lr_scheduler`` has been overwritten by ``learning_rate`` in optimizer.
2%|██▌                                                                                                                              | 1/50 [00:00<00:06,  7.65it/s, epoch=1/3, avg_epoch_loss=2.79]Epoch[0] gave nan loss
2%|██▌                                                                                                                              | 1/50 [00:00<00:11,  4.15it/s, epoch=1/3, avg_epoch_loss=2.79]
Traceback (most recent call last):
File "example_deepar_fails_mixture_learning.py", line 65, in <module>
predictor = estimator.train(training_data=data)
File "/Users/eeehrlic/Documents/projects/capstone/gluon-ts/src/gluonts/model/estimator.py", line 270, in train
return self.train_model(
File "/Users/eeehrlic/Documents/projects/capstone/gluon-ts/src/gluonts/model/estimator.py", line 245, in train_model
self.trainer(
File "/Users/eeehrlic/Documents/projects/capstone/gluon-ts/src/gluonts/mx/trainer/_base.py", line 376, in __call__
raise GluonTSUserError(
gluonts.core.exception.GluonTSUserError: Got NaN in first epoch. Try reducing initial learning rate.

# Fixed code
def log_prob(self, x: Tensor) -> Tensor:
    F = self.F
    alpha, beta = self.alpha, self.beta

    def gamma_log_prob(x, alpha, beta):
        return (
            alpha * F.log(beta) - F.gammaln(alpha) + (alpha - 1) * F.log(x) - beta * x
        )

    """
        The gamma_log_prob(x) above returns NaNs for x<=0. Wherever there are NaN in either of the F.where() conditional
        vectors, then F.where() returns NaN at that entry as well, due to its indicator function multiplication: 
        1*f(x) + np.nan*0 = nan, since np.nan*0 return nan. 
        Therefore replacing gamma_log_prob(x) with gamma_log_prob(abs(x) mitigates nan returns in cases of x<=0 without 
        altering the value in cases of x>0. 
        This is a known issue in pytorch as well https://github.com/pytorch/pytorch/issues/12986.
        """
    # mask zeros to prevent NaN gradients for x==0
    x_masked = F.where(x == 0, x.ones_like() * 0.5, x)

    return F.where(
        x > 0,
        gamma_log_prob(F.abs(x_masked), alpha, beta),
        -np.inf * F.ones_like(x),
    )