justindomke / pangolin

probabilistic programming focused on fun
GNU Affero General Public License v3.0
21 stars 1 forks source link

Ellipsis style indexing raises exception #7

Closed xidulu closed 6 months ago

xidulu commented 6 months ago
from pangolin.interface import (
    vmap, normal_scale, normal, RV, VMapDist,
    makerv, viz_upstream, exponential, bernoulli, print_upstream,
    Constant, plate, exp, Index
)

y = plate(N=7)(lambda:
            plate(N=5)(lambda:
               plate(N=6)(lambda:
                          normal_scale(0,1))))
print(y[[0, 1, 2], :, :].shape) # This works
print(y[[0, 1, 2], ...].shape) # This raises exception
justindomke commented 6 months ago

Experimental support for ellipsis installed: https://github.com/justindomke/pangolin/blob/main/pangolin/ir.py#L623

I haven't written systematic tests yet, but it seems to work for most cases and the code is in the end quite simple. Please let me know if you have any problems.