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

    # mask zeros for the Beta distribution input to prevent NaN gradients
    inputs = F.where(F.logical_or(x == 0, x == 1), x.zeros_like() + 0.5, x)

    # compute log density, case by case
    return F.where(
        x == 1,
        F.log(self.one_probability.broadcast_like(x)),
        F.where(
            x == 0,
            F.log(self.zero_probability.broadcast_like(x)),
            F.log(self.beta_probability) + self.beta_distribution.log_prob(inputs),
        ),
    )


# Traceback (AttributeError)
0%|          | 0/50 [00:00<?, ?it/s]
learning rate from ``lr_scheduler`` has been overwritten by ``learning_rate`` in optimizer.
0%|          | 0/50 [00:01<?, ?it/s]
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-24-94b4fecb46f1> in <module>
2     freq="W", prediction_length=4, trainer=trainer_1, distr_output=ZeroAndOneInflatedBetaOutput(),
3     context_length=26, num_cells=104, dropout_rate=0.1)
----> 4 predictor_4 = estimator_4.train(training_data=data_train_beta)

~/.conda/envs/cuda_test/lib/python3.6/site-packages/gluonts/model/estimator.py in train(self, training_data, validation_data, num_workers, num_prefetch, shuffle_buffer_length, **kwargs)
274             num_prefetch,
275             shuffle_buffer_length,
--> 276             **kwargs,
277         ).predictor

~/.conda/envs/cuda_test/lib/python3.6/site-packages/gluonts/model/estimator.py in train_model(self, training_data, validation_data, num_workers, num_prefetch, shuffle_buffer_length, **kwargs)
247             input_names=get_hybrid_forward_input_names(trained_net),
248             train_iter=training_data_loader,
--> 249             validation_iter=validation_data_loader,
250         )
251

~/.conda/envs/cuda_test/lib/python3.6/site-packages/gluonts/mx/trainer/_base.py in __call__(self, net, input_names, train_iter, validation_iter)
343                     )
344
--> 345                     epoch_loss = loop(epoch_no, train_iter)
346                     if is_validation_available:
347                         epoch_loss = loop(

~/.conda/envs/cuda_test/lib/python3.6/site-packages/gluonts/mx/trainer/_base.py in loop(epoch_no, batch_iter, is_training)
260                             if first_forward:
261                                 first_forward = False
--> 262                                 _ = net(*inputs)
263                                 if self.post_initialize_cb:
264                                     self.post_initialize_cb(net)

~/.conda/envs/cuda_test/lib/python3.6/site-packages/mxnet/gluon/block.py in __call__(self, *args)
680             hook(self, args)
681
--> 682         out = self.forward(*args)
683
684         for hook in self._forward_hooks.values():

~/.conda/envs/cuda_test/lib/python3.6/site-packages/mxnet/gluon/block.py in forward(self, x, *args)
1242                                      'Find all contexts = {}'.format(ctx_set))
1243                 with ctx:
-> 1244                     return self._call_cached_op(x, *args)
1245             with ctx:
1246                 try:

~/.conda/envs/cuda_test/lib/python3.6/site-packages/mxnet/gluon/block.py in _call_cached_op(self, *args)
993     def _call_cached_op(self, *args):
994         if self._cached_op is None:
--> 995             self._build_cache(*args)
996         assert self._cached_op, "Gluon failed to build the cache. " \
997                                 "This should never happen. " \

~/.conda/envs/cuda_test/lib/python3.6/site-packages/mxnet/gluon/block.py in _build_cache(self, *args)
922
923     def _build_cache(self, *args):
--> 924         data, out = self._get_graph(*args)
925         data_names = {data.name: i for i, data in enumerate(data)}
926         params = self.collect_params()

~/.conda/envs/cuda_test/lib/python3.6/site-packages/mxnet/gluon/block.py in _get_graph(self, *args)
914             params = {i: j.var() for i, j in self._reg_params.items()}
915             with self.name_scope():
--> 916                 out = self.hybrid_forward(symbol, *grouped_inputs, **params)  # pylint: disable=no-value-for-parameter
917             out, self._out_format = _flatten(out, "output")
918

~/.conda/envs/cuda_test/lib/python3.6/site-packages/gluonts/model/deepar/_network.py in hybrid_forward(self, F, feat_static_cat, feat_static_real, past_time_feat, past_target, past_observed_values, future_time_feat, future_target, future_observed_values)
456
457         # (batch_size, seq_len)
--> 458         loss = distr.loss(target)
459
460         # (batch_size, seq_len, *target_shape)

~/.conda/envs/cuda_test/lib/python3.6/site-packages/gluonts/mx/distribution/distribution.py in loss(self, x)
136             for each event in `x`.
137         """
--> 138         return -self.log_prob(x)
139
140     def prob(self, x: Tensor) -> Tensor:

~/.conda/envs/cuda_test/lib/python3.6/site-packages/gluonts/mx/distribution/transformed_distribution.py in log_prob(self, y)
147             y = x
148
--> 149         return self.base_distribution.log_prob(x) + lp
150
151     def cdf(self, y: Tensor) -> Tensor:

~/.conda/envs/cuda_test/lib/python3.6/site-packages/gluonts/mx/distribution/inflated_beta.py in log_prob(self, x)
80
81         # mask zeros for the Beta distribution input to prevent NaN gradients
---> 82         inputs = F.where(F.logical_or(x == 0, x == 1), x.zeros_like() + 0.5, x)
83
84         # compute log density, case by case

AttributeError: module 'mxnet.symbol' has no attribute 'logical_or'

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

    # mask zeros for the Beta distribution input to prevent NaN gradients
    inputs = F.where(F.broadcast_logical_or(x == 0, x == 1), x.zeros_like() + 0.5, x)

    # compute log density, case by case
    return F.where(
        x == 1,
        F.log(self.one_probability.broadcast_like(x)),
        F.where(
            x == 0,
            F.log(self.zero_probability.broadcast_like(x)),
            F.log(self.beta_probability) + self.beta_distribution.log_prob(inputs),
        ),
    )