code: def one_hot(x: TensorType, space: gym.Space) -> TensorType:
      """Returns a one-hot tensor, given and int tensor and a space.
  
      Handles the MultiDiscrete case as well.
  
      Args:
          x: The input tensor.
          space: The space to use for generating the one-hot tensor.
  
      Returns:
          The resulting one-hot tensor.
  
      Raises:
          ValueError: If the given space is not a discrete one.
  
      .. testcode::
  
          import torch
          import gymnasium as gym
          from ray.rllib.utils.torch_utils import one_hot
          x = torch.IntTensor([0, 3])  # batch-dim=2
          # Discrete space with 4 (one-hot) slots per batch item.
          s = gym.spaces.Discrete(4)
          print(one_hot(x, s))
          x = torch.IntTensor([[0, 1, 2, 3]])  # batch-dim=1
          # MultiDiscrete space with 5 + 4 + 4 + 7 = 20 (one-hot) slots
          # per batch item.
          s = gym.spaces.MultiDiscrete([5, 4, 4, 7])
          print(one_hot(x, s))
  
      .. testoutput::
  
          tensor([[1, 0, 0, 0],
                  [0, 0, 0, 1]])
          tensor([[1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0]])
      """
      if isinstance(space, Discrete):
          return nn.functional.one_hot(x.long(), space.n)
      elif isinstance(space, MultiDiscrete):
          if isinstance(space.nvec[0], np.ndarray):
              nvec = np.ravel(space.nvec)
              x = x.reshape(x.shape[0], -1)
          else:
              nvec = space.nvec
          return torch.cat(
              [
                  nn.functional.one_hot(x[:, i].long(), n)
                  for i, n in enumerate(nvec)
              ],
              dim=-1,
          )
      else:
          raise ValueError(
              "Unsupported space for `one_hot`: {}".format(space)
          )
  
imports: import logging
  import os
  import warnings
  from typing import Dict, List, Optional, TYPE_CHECKING, Union
  import gymnasium
  import numpy
  import tree
  from gymnasium.spaces import Discrete, MultiDiscrete
  from packaging import version
  import ray
  from ray.rllib.models.repeated_values import RepeatedValues
  from ray.rllib.utils.annotations import Deprecated, PublicAPI, DeveloperAPI
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.numpy import SMALL_NUMBER
  from ray.rllib.utils.typing import LocalOptimizer, SpaceStruct, TensorStructType, TensorType
label_desc: Use image node_id 15 for calling a global function with example usage: one_hot(x, space) and returns: nn, torch