cc @fritzo (sorry, didn't notice you're already tagged)

cc @neerajprad

It seems a little forceful to always enable grad. One option that seems consistent to me is to inherit grad_enabled from its setting at initialization time

class Distribution(object):
    def __init__(self, ...):
        ...
        self._grad_enabled = torch.autograd.get_grad_enabled()  # what is the actual syntax?

class lazy_property(object):
    ...
    def __get__(self, instance, obj_type=None):
        if instance is None:
            return self
        with torch.autograd.set_grad_enabled(instance._grad_enabled):
            value = self.wrapped(instance)
        setattr(instance, self.wrapped.__name__, value)
        return value

However it's unclear whether we should be clever or whether we should simply require strict usage: "grad_enabled should have a single value throughout the lifetime of a distribution object". The distributions are intended to be flyweight objects and to be cheap to reconstruct in each grad_enabled context.

I think the "constant grad mode" invariant is a bit too easy to get wrong. Forcing grad mode doesn't seem that bad really. If your distribution parameters don't require grad then it will be a no-op. Otherwise it's quite likely that you will be interested in differentiating those parts, and I think it's better to trade off some memory for compute in this case.

It's a hard problem, because you don't have any information about how the object will be used. You either need to drop the autograd history, hoping that someone who calls .sample() will never need it, or you need to just live with the fact that you might end up wasting some memory, because we temporarily enabled grad. Provided the distribution is short-lived the downside is somewhat irrelevant though, because the properties should get removed quickly. If we find out that this strategy doesn't work too well, we can always add an extra constructor parameter that lets you choose a strategy for caching those.

@fritzo @apaszke @neerajprad have we reached consensus on whether we should merge this? :)

This seems reasonable to me (I'm convinced by @apaszke).