set_composite_lp_aggregate#

class tensordict.nn.set_composite_lp_aggregate(mode: bool = True, *, thread_local: bool = False)#

Bases: _DecoratorContextManager

Controls whether CompositeDistribution log-probabilities and entropies will be aggregated in a single tensor.

When composite_lp_aggregate() returns True, the log-probs / entropies of CompositeDistribution will be summed into a single tensor with the shape of the root tensordict. When it returns False (the default), they are returned as a TensorDict with one <name>_log_prob (or <name>_entropy) entry per sampled key. The value also decides the keys where probabilistic modules write their log-probabilities: see the log-probability key table of ProbabilisticTensorDictModule, and Modules: TensorDictModule and tensordict.nn.

The value of composite_lp_aggregate can also be controlled through the COMPOSITE_LP_AGGREGATE environment variable. It works as a context manager and as a decorator, and set() and unset() change it outside of a block.

Parameters:: mode (bool, optional) – the value that composite_lp_aggregate() returns inside the block. Defaults to True.
Keyword Arguments:: thread_local (bool, optional) – if True, the mode applies to the calling thread only, and blocks that other threads enter do not change it. Inside the block, blocks without this flag that the same thread enters stay in that thread too. Threads started inside the block see the process-wide mode. If False, the mode applies to the whole process. Defaults to False.

Example

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import CompositeDistribution, set_composite_lp_aggregate
>>> from torch import distributions as d
>>> _ = torch.manual_seed(0)
>>> params = TensorDict({
...     "cont": {"loc": torch.randn(3, 4), "scale": torch.rand(3, 4)},
...     ("nested", "disc"): {"logits": torch.randn(3, 10)}
... }, [3])
>>> dist = CompositeDistribution(params,
...     distribution_map={"cont": d.Normal, ("nested", "disc"): d.Categorical})
>>> sample = dist.sample((4,))
>>> with set_composite_lp_aggregate(False):
...     lp = dist.log_prob(sample)
...     print(lp)
TensorDict(
    fields={
        cont_log_prob: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                disc_log_prob: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4, 3]),
    device=None,
    is_shared=False)
>>> with set_composite_lp_aggregate(True):
...     lp = dist.log_prob(sample)
...     print(lp)
tensor([[-2.0886, -1.2155, -0.0414],
        [-2.8973, -5.5165,  2.4402],
        [-0.2806, -1.2799,  3.1733],
        [-3.0407, -4.3593,  0.5763]])

set_composite_lp_aggregate#

Docs

Tutorials

Resources