code: def spatial_softmax2d(
      input: Tensor, temperature: Tensor = torch.tensor(1.0)
  ) -> Tensor:
      r"""Apply the Softmax function over features in each image channel.
  
      Note that this function behaves differently to :py:class:`torch.nn.Softmax2d`, which
      instead applies Softmax over features at each spatial location.
  
      Args:
          input: the input tensor with shape :math:`(B, N, H, W)`.
          temperature: factor to apply to input, adjusting the "smoothness" of the output distribution.
  
      Returns:
         a 2D probability distribution per image channel with shape :math:`(B, N, H, W)`.
  
      Examples:
          >>> heatmaps = torch.tensor([[[
          ... [0., 0., 0.],
          ... [0., 0., 0.],
          ... [0., 1., 2.]]]])
          >>> spatial_softmax2d(heatmaps)
          tensor([[[[0.0585, 0.0585, 0.0585],
                    [0.0585, 0.0585, 0.0585],
                    [0.0585, 0.1589, 0.4319]]]])
      """
      _validate_batched_image_tensor_input(input)
  
      batch_size, channels, height, width = input.shape
      temperature = temperature.to(
          device=input.device, dtype=input.dtype
      )
      x = input.view(batch_size, channels, -1)
  
      x_soft = softmax(x * temperature, dim=-1)
  
      return x_soft.view(batch_size, channels, height, width)
  
file_path: kornia/kornia/geometry/subpix/dsnt.py
imports: from __future__ import annotations
  import torch
  from kornia.core import Tensor, concatenate, softmax
  from kornia.core.check import KORNIA_CHECK_IS_TENSOR, KORNIA_CHECK_SHAPE
  from kornia.utils.grid import create_meshgrid
label_desc: Use image node_id 2 for calling a global function with example usage: spatial_softmax2d(input, temperature) and returns: x_soft