Implement sizesof for named and positional distributions

[jfeser]

Closes #213

[eb8680]

I think it would be better and not that much more involved to address the root cause #203 instead of special-casing handlers.torch.sizesof

[jfeser]

Ok, I'll rework this into something that addresses #203

[jfeser]

Handling NamedDistribution is tricky, because the named dimensions don't appear as arguments to torch_getitem. We would need to switch from wrapping the distribution to reconstructing it with named arguments. For most (all?) distributions, naming the batch dimensions of the parameters is the same as naming the batch dimensions of the sample results, so this approach should still work. It's certainly more involved than the original version of this PR, because we need to introspect the distributions and reconstruct them, which might need per-distribution logic.

[eb8680]

This mostly makes sense, I just have a few comments

[eb8680]

Hmm, why does positional_distribution needs to be a separate operation, as opposed to just using to_tensor on defterm(d)?

[jfeser]

We could move this logic into to_tensor, sure. The default behavior of to_tensor doesn't work though, because vmap is restricted to only return tensors.

[jfeser]

As discussed, I'd like to leave this alone for now.

[eb8680]

Looks good, thanks for taking a pass at this! I think there are still a couple of easy simplifications of positional_distribution and named_distribution as noted in comments, but you can address them in a followup PR if this is blocking other work.

[eb8680]

CI failure is just a randomized distribution test.

[eb8680]

This should be registered with and accessed through defdata, but we can also address that in a followup PR.