code: def softmax(x, axis=-1):
      x = convert_to_tensor(x)
      dtype = standardize_dtype(x.dtype)
      # TODO: tnn.softmax doesn't support float16 using cpu
      if (
          get_device() == "cpu"
          and standardize_dtype(x.dtype) == "float16"
      ):
          x = cast(x, "float32")
      if axis is None:
          # Unlike numpy, PyTorch will handle axis=None as axis=-1.
          # We need this workaround for the reduction on every dim.
          output = torch.reshape(x, [-1])
          output = tnn.softmax(output, dim=-1)
          output = torch.reshape(output, x.shape)
      else:
          output = tnn.softmax(x, dim=axis)
      return cast(output, dtype)
  
imports: import torch
  import torch.nn.functional
  import tree
  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_torch
  from keras.backend.config import epsilon
  from keras.backend.torch.core import cast
  from keras.backend.torch.core import convert_to_tensor
  from keras.backend.torch.core import get_device
  from keras.backend.torch.numpy import expand_dims
  from keras.backend.torch.numpy import maximum
  from keras.backend.torch.numpy import where
  from keras.utils.argument_validation import standardize_tuple
label_desc: Use image node_id 15 for calling a global function with example usage: softmax(x, axis) and returns: cast