code: def one_hot(x, num_classes, axis=-1, dtype="float32"):
      x = convert_to_tensor(x)
      input_shape = x.shape
  
      # Shrink the last dimension if the shape is (..., 1).
      if input_shape and input_shape[-1] == 1 and len(input_shape) > 1:
          input_shape = tuple(input_shape[:-1])
  
      x = x.reshape(-1)
      if not num_classes:
          num_classes = np.max(x) + 1
  
      batch_size = x.shape[0]
      categorical = np.zeros((batch_size, num_classes), dtype=dtype)
      valid_indices = x >= 0
      categorical[
          np.arange(batch_size)[valid_indices], x[valid_indices]
      ] = 1
  
      # First, reshape the array with the extra dimension at the end
      output_shape = input_shape + (num_classes,)
      categorical = np.reshape(categorical, output_shape)
  
      # Then, move this new dimension to the right place (according to axis)
      if axis != -1:
          categorical = np.moveaxis(categorical, -1, axis)
  
      return categorical
  
imports: import jax
  import numpy
  from jax import lax
  from jax import numpy
  from keras.backend import standardize_data_format
  from keras.backend import standardize_dtype
  from keras.backend.common.backend_utils import compute_conv_transpose_padding_args_for_jax
  from keras.backend.config import epsilon
  from keras.backend.numpy.core import cast
  from keras.backend.numpy.core import convert_to_tensor
  from keras.backend.numpy.core import is_tensor
  from keras.utils.module_utils import scipy
label_desc: Use image node_id 26 for calling a global function with example usage: one_hot(x, num_classes, axis, dtype) and returns: categorical