From 097d8ad9879f8e83c89c7d8da0a4727146355cfd Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Sun, 24 Nov 2024 09:13:36 +0100 Subject: [PATCH] [Feature] spec.is_empty(recurse) ghstack-source-id: faa3b1df5133c77462d6dd013d3854d684cc7e94 Pull Request resolved: https://github.com/pytorch/rl/pull/2596 --- torchrl/data/tensor_specs.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py index 2ef74bb4521..1f31db01ec7 100644 --- a/torchrl/data/tensor_specs.py +++ b/torchrl/data/tensor_specs.py @@ -4278,8 +4278,22 @@ def shape(self, value: torch.Size): ) self._shape = _size(value) - def is_empty(self): - """Whether the composite spec contains specs or not.""" + def is_empty(self, recurse: bool = False): + """Whether the composite spec contains specs or not. + + Args: + recurse (bool): whether to recursively assess if the spec is empty. + If ``True``, will return ``True`` if there are no leaves. If ``False`` + (default) will return whether there is any spec defined at the root level. + + """ + if recurse: + for spec in self._specs.values(): + if spec is None: + continue + if isinstance(spec, Composite) and spec.is_empty(recurse=True): + continue + return False return len(self._specs) == 0 @property