def _base_dist(self, p: TensorLike, *args, **kwargs): return tfd.Geometric(probs=p, *args, **kwargs)
def _init_distribution(conditions): p = conditions["p"] return tfd.Geometric(probs=p)
def _init_distribution(conditions, **kwargs): probs = conditions["probs"] return tfd.Geometric(probs=probs, **kwargs)