Skip to content

Conversation

ColtAllen
Copy link

Closes #573

test_random_edge_cases uses rather extreme values and may be unstable.

Copy link
Member

@ricardoV94 ricardoV94 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I mentioned in the original issue that I think this makes more sense in pymc-marketing directly. Just left one comment that is relevant regardless of where you implement it in the end

Comment on lines 405 to 424
class ShiftedBetaGeometricRV(RandomVariable):
name = "sbg"
signature = "(),()->()"

dtype = "int64"
_print_name = ("ShiftedBetaGeometric", "\\operatorname{ShiftedBetaGeometric}")

@classmethod
def rng_fn(cls, rng, alpha, beta, size):
if size is None:
size = np.broadcast_shapes(alpha.shape, beta.shape)

alpha = np.broadcast_to(alpha, size)
beta = np.broadcast_to(beta, size)

p = rng.beta(a=alpha, b=beta, size=size)

samples = rng.geometric(p, size=size)

return samples
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Use a SymbolicRandomVariable that defines the random method in terms of pytensor operations. That way it will work automatically in all backends (jax/numba), not just C/python

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add ShiftedBetaGeometric Distribution
2 participants