code: def conv_soft_argmax3d(
      input: Tensor,
      kernel_size: tuple[int, int, int] = (3, 3, 3),
      stride: tuple[int, int, int] = (1, 1, 1),
      padding: tuple[int, int, int] = (1, 1, 1),
      temperature: Tensor | float = tensor(1.0),
      normalized_coordinates: bool = False,
      eps: float = 1e-8,
      output_value: bool = True,
      strict_maxima_bonus: float = 0.0,
  ) -> Tensor | tuple[Tensor, Tensor]:
      r"""Compute the convolutional spatial Soft-Argmax 3D over the windows of a given heatmap.
  
      .. math::
               ijk(X) = \frac{\sum{(i,j,k)} * exp(x / T)  \in X} {\sum{exp(x / T)  \in X}}
  
      .. math::
               val(X) = \frac{\sum{x * exp(x / T)  \in X}} {\sum{exp(x / T)  \in X}}
  
      where ``T`` is temperature.
  
      Args:
          input: the given heatmap with shape :math:`(N, C, D_{in}, H_{in}, W_{in})`.
          kernel_size:  size of the window.
          stride: stride of the window.
          padding: input zero padding.
          temperature: factor to apply to input.
          normalized_coordinates: whether to return the coordinates normalized in the range of :math:[-1, 1]`.
              Otherwise, it will return the coordinates in the range of the input shape.
          eps: small value to avoid zero division.
          output_value: if True, val is output, if False, only ij.
          strict_maxima_bonus: pixels, which are strict maxima will score (1 + strict_maxima_bonus) * value.
            This is needed for mimic behavior of strict NMS in classic local features
  
      Returns:
          Function has two outputs - argmax coordinates and the softmaxpooled heatmap values themselves.
          On each window, the function computed returns with shapes :math:`(N, C, 3, D_{out}, H_{out}, W_{out})`,
          :math:`(N, C, D_{out}, H_{out}, W_{out})`,
  
          where
  
           .. math::
               D_{out} = \left\lfloor\frac{D_{in}  + 2 \times \text{padding}[0] -
               (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor
  
           .. math::
               H_{out} = \left\lfloor\frac{H_{in}  + 2 \times \text{padding}[1] -
               (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor
  
           .. math::
               W_{out} = \left\lfloor\frac{W_{in}  + 2 \times \text{padding}[2] -
               (\text{kernel\_size}[2] - 1) - 1}{\text{stride}[2]} + 1\right\rfloor
  
      Examples:
          >>> input = torch.randn(20, 16, 3, 50, 32)
          >>> nms_coords, nms_val = conv_soft_argmax3d(input, (3, 3, 3), (1, 2, 2), (0, 1, 1))
      """
      if not torch.is_tensor(input):
          raise TypeError(
              f"Input type is not a Tensor. Got {type(input)}"
          )
  
      if not len(input.shape) == 5:
          raise ValueError(
              f"Invalid input shape, we expect BxCxDxHxW. Got: {input.shape}"
          )
  
      if temperature <= 0:
          raise ValueError(
              f"Temperature should be positive float or tensor. Got: {temperature}"
          )
  
      b, c, d, h, w = input.shape
      kz, ky, kx = kernel_size
      device: torch.device = input.device
      dtype: torch.dtype = input.dtype
      input = input.view(b * c, 1, d, h, w)
  
      center_kernel: Tensor = _get_center_kernel3d(
          kz, ky, kx, device
      ).to(dtype)
      window_kernel: Tensor = _get_window_grid_kernel3d(
          kz, ky, kx, device
      ).to(dtype)
  
      # applies exponential normalization trick
      # https://timvieira.github.io/blog/post/2014/02/11/exp-normalize-trick/
      # https://github.com/pytorch/pytorch/blob/bcb0bb7e0e03b386ad837015faba6b4b16e3bfb9/aten/src/ATen/native/SoftMax.cpp#L44
      x_max = F.adaptive_max_pool3d(input, (1, 1, 1))
  
      # max is detached to prevent undesired backprop loops in the graph
      x_exp = ((input - x_max.detach()) / temperature).exp()
  
      pool_coef: float = float(kx * ky * kz)
  
      # softmax denominator
      den = (
          pool_coef
          * F.avg_pool3d(
              x_exp.view_as(input),
              kernel_size,
              stride=stride,
              padding=padding,
          )
          + eps
      )
  
      # We need to output also coordinates
      # Pooled window center coordinates
      grid_global: Tensor = (
          create_meshgrid3d(d, h, w, False, device=device)
          .to(dtype)
          .permute(0, 4, 1, 2, 3)
      )
  
      grid_global_pooled = F.conv3d(
          grid_global, center_kernel, stride=stride, padding=padding
      )
  
      # Coordinates of maxima residual to window center
      # prepare kernel
      coords_max: Tensor = F.conv3d(
          x_exp, window_kernel, stride=stride, padding=padding
      )
  
      coords_max = coords_max / den.expand_as(coords_max)
      coords_max = coords_max + grid_global_pooled.expand_as(coords_max)
      # [:,:, 0, ...] is depth (scale)
      # [:,:, 1, ...] is x
      # [:,:, 2, ...] is y
  
      if normalized_coordinates:
          coords_max = normalize_pixel_coordinates3d(
              coords_max.permute(0, 2, 3, 4, 1), d, h, w
          )
          coords_max = coords_max.permute(0, 4, 1, 2, 3)
  
      # Back B*C -> (b, c)
      coords_max = coords_max.view(
          b,
          c,
          3,
          coords_max.size(2),
          coords_max.size(3),
          coords_max.size(4),
      )
  
      if not output_value:
          return coords_max
  
      x_softmaxpool = (
          pool_coef
          * F.avg_pool3d(
              x_exp.view(input.size()) * input,
              kernel_size,
              stride=stride,
              padding=padding,
          )
          / den
      )
      if strict_maxima_bonus > 0:
          in_levels: int = input.size(2)
          out_levels: int = x_softmaxpool.size(2)
          skip_levels: int = (in_levels - out_levels) // 2
          strict_maxima: Tensor = F.avg_pool3d(
              nms3d(input, kernel_size), 1, stride, 0
          )
          strict_maxima = strict_maxima[
              :, :, skip_levels : out_levels - skip_levels
          ]
          x_softmaxpool *= 1.0 + strict_maxima_bonus * strict_maxima
      x_softmaxpool = x_softmaxpool.view(
          b,
          c,
          x_softmaxpool.size(2),
          x_softmaxpool.size(3),
          x_softmaxpool.size(4),
      )
      return coords_max, x_softmaxpool
  
imports: from __future__ import annotations
  import torch
  import torch.nn.functional
  from kornia.core import Module, Tensor, concatenate, rand, stack, tensor, where, zeros, zeros_like
  from kornia.filters.sobel import spatial_gradient3d
  from kornia.geometry.conversions import normalize_pixel_coordinates, normalize_pixel_coordinates3d
  from kornia.utils import create_meshgrid, create_meshgrid3d
  from kornia.utils._compat import torch_version_ge
  from kornia.utils.helpers import safe_solve_with_mask
  from .dsnt import spatial_expectation2d, spatial_softmax2d
  from .nms import nms3d
label_desc: Use image node_id 6 for calling a global function with example usage: conv_soft_argmax3d(input, kernel_size, stride, padding, temperature, normalized_coordinates, eps, output_value, strict_maxima_bonus) and returns: coords_max,  x_softmaxpool, coords_max