pyro-ppl / funsor

Functional tensors for probabilistic programming
https://funsor.pyro.ai
Apache License 2.0
236 stars 20 forks source link

Allow stack Numbers #506

Closed fehiepsi closed 3 years ago

fehiepsi commented 3 years ago

Fixes https://github.com/pyro-ppl/numpyro/issues/964

After https://github.com/pyro-ppl/funsor/pull/491, we are no longer able to stack Numbers. This PR relaxes this restriction because numpy, jax backends can perform stack over Python scalars

import funsor; funsor.set_backend("jax")
from funsor import ops, Number
ops.stack((Number(2.), Number(3)))