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::
          :skipif: True
  
          import gymnasium as gym
          import tensorflow as tf
          from ray.rllib.utils.tf_utils import one_hot
          x = tf.Variable([0, 3], dtype=tf.int32)  # batch-dim=2
          # Discrete space with 4 (one-hot) slots per batch item.
          s = gym.spaces.Discrete(4)
          one_hot(x, s)
  
      .. testoutput::
  
          <tf.Tensor 'one_hot:0' shape=(2, 4) dtype=float32>
  
      .. testcode::
          :skipif: True
  
          x = tf.Variable([[0, 1, 2, 3]], dtype=tf.int32)  # 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])
          one_hot(x, s)
  
      .. testoutput::
  
          <tf.Tensor 'concat:0' shape=(1, 20) dtype=float32>
      """
      if isinstance(space, Discrete):
          return tf.one_hot(x, space.n, dtype=tf.float32)
      elif isinstance(space, MultiDiscrete):
          if isinstance(space.nvec[0], np.ndarray):
              nvec = np.ravel(space.nvec)
              x = tf.reshape(x, (x.shape[0], -1))
          else:
              nvec = space.nvec
          return tf.concat(
              [
                  tf.one_hot(x[:, i], n, dtype=tf.float32)
                  for i, n in enumerate(nvec)
              ],
              axis=-1,
          )
      else:
          raise ValueError(
              "Unsupported space for `one_hot`: {}".format(space)
          )
  
imports: import logging
  from typing import Any, Callable, List, Optional, Type, TYPE_CHECKING, Union
  import gymnasium
  import numpy
  import tree
  from gymnasium.spaces import Discrete, MultiDiscrete
  from ray.rllib.utils.annotations import PublicAPI, DeveloperAPI
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space
  from ray.rllib.utils.typing import LocalOptimizer, ModelGradients, PartialAlgorithmConfigDict, SpaceStruct, TensorStructType, TensorType
label_desc: Use image node_id 11 for calling a global function with example usage: one_hot(x, space) and returns: tf, tf