class_docstr: Action distribution that operates on multiple, possibly nested actions.
class_docstr_tok: ["Action","distribution","that","operates","on","multiple",",","possibly","nested","actions","."]
code: def logp(self, value: TensorType) -> TensorType:
      # Different places in RLlib use this method with different inputs.
      # We therefore need to handle a flattened and concatenated input, as well as
      # a nested one.
      # TODO(Artur): Deprecate tensor inputs, only allow nested structures.
      if isinstance(value, torch.Tensor):
          split_indices = []
          for dist in self._flat_child_distributions:
              if isinstance(dist, TorchCategorical):
                  split_indices.append(1)
              elif isinstance(dist, TorchMultiCategorical):
                  split_indices.append(len(dist._cats))
              else:
                  sample = dist.sample()
                  # Cover Box(shape=()) case.
                  if len(sample.shape) == 1:
                      split_indices.append(1)
                  else:
                      split_indices.append(sample.size()[1])
          split_value = list(torch.split(value, split_indices, dim=1))
      else:
          split_value = tree.flatten(value)
  
      def map_(val, dist):
          # Remove extra dimension if present.
          if (
              isinstance(dist, TorchCategorical)
              and val.shape[-1] == 1
              and len(val.shape) > 1
          ):
              val = torch.squeeze(val, dim=-1)
          return dist.logp(val)
  
      flat_logps = tree.map_structure(
          map_, split_value, self._flat_child_distributions
      )
  
      return sum(flat_logps)
  
imports: import gymnasium
  import numpy
  from typing import Optional, List, Mapping, Iterable, Dict
  import tree
  import abc
  from ray.rllib.models.distributions import Distribution
  from ray.rllib.utils.annotations import override, DeveloperAPI
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.typing import TensorType, Union, Tuple
label_desc: Use image node_id 3 for calling the TorchMultiDistribution obj's underlying member method code with example usage: obj.logp(value) and returns: sum, dist