broadcast numerical types to jnp.array for flexible shapes

[Qazalbash]

Description

rvs = Beta(alpha=1.1, beta=5.5).rvs(10) # fine
rvs = Beta(alpha=jnp.array([1.1, 1.0]), beta=jnp.array([5.5, 3.0])).rvs(10) # gives error.

error is,

ValueError: beta parameter shapes must be broadcast-compatible with shape argument, and the result of broadcasting the shapes must equal the shape argument, but got result (10, 2) for shape argument (10, 1).

Proposal

It would be great to write a generic function that broadcasts every element to an array and define it in some general folder like utils.