class_docstr: Wrapper class for PyTorch Normal distribution.
  
  Creates a normal distribution parameterized by :attr:`loc` and :attr:`scale`. In
  case of multi-dimensional distribution, the variance is assumed to be diagonal.
  
  .. testcode::
      :skipif: True
  
      loc, scale = torch.tensor([0.0, 0.0]), torch.tensor([1.0, 1.0])
      m = TorchDiagGaussian(loc=loc, scale=scale)
      m.sample(sample_shape=(2,))  # 2d normal dist with loc=0 and scale=1
  
  .. testoutput::
  
      tensor([[ 0.1046, -0.6120], [ 0.234, 0.556]])
  
  .. testcode::
      :skipif: True
  
      # scale is None
      m = TorchDiagGaussian(loc=torch.tensor([0.0, 1.0]))
      m.sample(sample_shape=(2,))  # normally distributed with loc=0 and scale=1
  
  .. testoutput::
  
      tensor([0.1046, 0.6120])
  
  
  Args:
      loc: mean of the distribution (often referred to as mu). If scale is None, the
          second half of the `loc` will be used as the log of scale.
      scale: standard deviation of the distribution (often referred to as sigma).
          Has to be positive.
class_docstr_tok: ["Wrapper","class","for","PyTorch","Normal","distribution",".","Creates","a","normal","distribution","parameterized","by",":","attr",":","`","loc","`","and",":","attr",":","`","scale","`",".","In","case","of","multi-dimensional","distribution",",","the","variance","is","assumed","to","be","diagonal",".","..","testcode",":",":",":","skipif",":","True","loc",",","scale","=","torch.tensor","(","[","0.0",",","0.0","]",")",",","torch.tensor","(","[","1.0",",","1.0","]",")","m","=","TorchDiagGaussian","(","loc=loc",",","scale=scale",")","m.sample","(","sample_shape=","(","2",",",")",")","#","2d","normal","dist","with","loc=0","and","scale=1","..","testoutput",":",":","tensor","(","[","[","0.1046",",","-0.6120","]",",","[","0.234",",","0.556","]","]",")","..","testcode",":",":",":","skipif",":","True","#","scale","is","None","m","=","TorchDiagGaussian","(","loc=torch.tensor","(","[","0.0",",","1.0","]",")",")","m.sample","(","sample_shape=","(","2",",",")",")","#","normally","distributed","with","loc=0","and","scale=1","..","testoutput",":",":","tensor","(","[","0.1046",",","0.6120","]",")","Args",":","loc",":","mean","of","the","distribution","(","often","referred","to","as","mu",")",".","If","scale","is","None",",","the","second","half","of","the","`","loc","`","will","be","used","as","the","log","of","scale",".","scale",":","standard","deviation","of","the","distribution","(","often","referred","to","as","sigma",")",".","Has","to","be","positive","."]
code: def kl(self, other: "TorchDistribution") -> TensorType:
      return super().kl(other).sum(-1)
  
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 5 for calling the TorchDiagGaussian obj's underlying member method code with example usage: obj.kl(other) and returns: super