code: def _elementary_basis_index(axis: str) -> int:
      if axis == "x":
          return 0
      elif axis == "y":
          return 1
      elif axis == "z":
          return 2
      raise ValueError(
          f"Expected axis to be from ['x', 'y', 'z'], got {axis}"
      )
  
imports: from __future__ import annotations
  import functools
  import re
  import typing
  import scipy.spatial.transform
  import jax
  import jax.numpy
  from jax._src.numpy.util import _wraps
label_desc: Use image node_id 7 for calling a global function with example usage: _elementary_basis_index(axis) and returns: int, int, int