parameters: (key, shape)
return_statement: return expanded_key, prepend
docstring: Make implicit axes explicit by adding ``slice(None)``
      and convert boolean array to integer array through `nonzero`.
  
      Examples
      --------
      >>> shape = (3, 4, 5)
      >>> indexing_key_expand_implicit_axes(np.s_[2, 1, 1], shape)
      (2, 1, 1)
      >>> indexing_key_expand_implicit_axes(np.s_[0], shape)
      (0, slice(None, None, None), slice(None, None, None))
      >>> indexing_key_expand_implicit_axes(np.s_[0, ...], shape)  # equivalent
      (0, slice(None, None, None), slice(None, None, None))
      >>> indexing_key_expand_implicit_axes(np.s_[:2, None, 0, ...], shape)
      (slice(None, 2, None), None, 0, slice(None, None, None))
      >>> bool_array = np.array([[True, False, True, False],
                                 [False, True, False, True],
                                 [True, False, True, False]], dtype=np.bool)
      >>> indexing_key_expand_implicit_axes(np.s_[bool_array, None, 0:2], shape)
      (array([0, 0, 1, 1, 2, 2], dtype=int64), array([0, 2, 1, 3, 0, 2], dtype=int64), None, slice(None, 2, None))
docstring_summary: Make implicit axes explicit by adding ``slice(None)``
      and convert boolean array to integer array through `nonzero`.
docstring_tokens: Make implicit axes explicit by adding slice ( None ) and convert boolean array to integer array through nonzero .
function: def indexing_key_expand_implicit_axes(key, shape):
      """
      Make implicit axes explicit by adding ``slice(None)``
      and convert boolean array to integer array through `nonzero`.
  
      Examples
      --------
      >>> shape = (3, 4, 5)
      >>> indexing_key_expand_implicit_axes(np.s_[2, 1, 1], shape)
      (2, 1, 1)
      >>> indexing_key_expand_implicit_axes(np.s_[0], shape)
      (0, slice(None, None, None), slice(None, None, None))
      >>> indexing_key_expand_implicit_axes(np.s_[0, ...], shape)  # equivalent
      (0, slice(None, None, None), slice(None, None, None))
      >>> indexing_key_expand_implicit_axes(np.s_[:2, None, 0, ...], shape)
      (slice(None, 2, None), None, 0, slice(None, None, None))
      >>> bool_array = np.array([[True, False, True, False],
                                 [False, True, False, True],
                                 [True, False, True, False]], dtype=np.bool)
      >>> indexing_key_expand_implicit_axes(np.s_[bool_array, None, 0:2], shape)
      (array([0, 0, 1, 1, 2, 2], dtype=int64), array([0, 2, 1, 3, 0, 2], dtype=int64), None, slice(None, 2, None))
      """
      if not isinstance(key, tuple):
          key = (key,)
      # We need to loop explicitly since tuple functions like `index()` or
      # `count()` use `==` internally, which doesn't play well with fancy
      # indexing.
      ell_idx = None
      num_none = 0
      nonell_key = []
  
      # For 0-d boolean indices: A new axis is added,
      # but at the same time no axis is "used". So if we have True,
      # we add a new axis (a bit like with np.newaxis). If it is
      # False, we add a new axis, but this axis has 0 entries.
      # prepend is defined to handle this case.
      # prepend = _NDARRAY_NO_ZERO_DIM_BOOL_ARRAY/-1 means there is no 0-d boolean scalar
      # prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE/0 means an zero dim must be expanded
      # prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE/1 means a new axis must be expanded
      prepend = _NDARRAY_NO_ZERO_DIM_BOOL_ARRAY
      axis = 0
      for i, idx in enumerate(key):
          if idx is Ellipsis:
              if ell_idx is not None:
                  raise IndexError(
                      'Cannot use more than one ellipsis (`...`) for indexing'
                  )
              ell_idx = i
          else:
              # convert primitive type boolean value to mx.np.bool type
              # otherwise will be treated as 1/0
              if isinstance(idx, bool):
                  idx = array(idx, dtype=np.bool_)
              if idx is None:
                  num_none += 1
              if isinstance(idx, NDArrayBase) and idx.ndim == 0 and idx.dtype == np.bool_:
                  if not idx: # array(False) has priority
                      prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE
                  else:
                      prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE
              elif isinstance(idx, NDArrayBase) and idx.ndim == 0 and idx.dtype != np.bool_:
                  # This handles ndarray of zero dim. e.g array(1)
                  # while advoid converting zero dim boolean array
                  # float type will be converted to int
                  nonell_key.append(int(idx.item()))
                  axis += 1
              elif isinstance(idx, NDArrayBase) and idx.dtype == np.bool_:
                  # Necessary size check before using `nonzero`
                  check_boolean_array_dimension(shape, axis, idx.shape)
                  # If the whole array is false and npx.set_np() is not set_up
                  # the program will throw infer shape error
                  if not is_np_array():
                      raise ValueError('Cannot perform boolean indexing in legacy mode. Please activate'
                                       ' numpy semantics by calling `npx.set_np()` in the global scope'
                                       ' before calling this function.')
                  # Add the arrays from the nonzero result to the index
                  nonell_key.extend(idx.nonzero())
                  axis += idx.ndim
              else:
                  nonell_key.append(idx)
                  axis += 1
  
      nonell_key = tuple(nonell_key)
  
      if ell_idx is None:
          # This handles the case of "too few" indices, e.g., `nd.zeros((2, 3))[0]`,
          # where the ellipsis is implicitly after the last entry.
          ell_idx = len(nonell_key)
  
      ell_ndim = len(shape) + num_none - len(nonell_key)
      expanded_key = (nonell_key[:ell_idx] +
                      (slice(None),) * ell_ndim +
                      nonell_key[ell_idx:])
  
      return expanded_key, prepend
function_tokens: def indexing_key_expand_implicit_axes ( key , shape ) : if not isinstance ( key , tuple ) : key = ( key , ) # We need to loop explicitly since tuple functions like `index()` or # `count()` use `==` internally, which doesn't play well with fancy # indexing. ell_idx = None num_none = 0 nonell_key = [ ] # For 0-d boolean indices: A new axis is added, # but at the same time no axis is "used". So if we have True, # we add a new axis (a bit like with np.newaxis). If it is # False, we add a new axis, but this axis has 0 entries. # prepend is defined to handle this case. # prepend = _NDARRAY_NO_ZERO_DIM_BOOL_ARRAY/-1 means there is no 0-d boolean scalar # prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE/0 means an zero dim must be expanded # prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE/1 means a new axis must be expanded prepend = _NDARRAY_NO_ZERO_DIM_BOOL_ARRAY axis = 0 for i , idx in enumerate ( key ) : if idx is Ellipsis : if ell_idx is not None : raise IndexError ( 'Cannot use more than one ellipsis (`...`) for indexing' ) ell_idx = i else : # convert primitive type boolean value to mx.np.bool type # otherwise will be treated as 1/0 if isinstance ( idx , bool ) : idx = array ( idx , dtype = np . bool_ ) if idx is None : num_none += 1 if isinstance ( idx , NDArrayBase ) and idx . ndim == 0 and idx . dtype == np . bool_ : if not idx : # array(False) has priority prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_FALSE else : prepend = _NDARRAY_ZERO_DIM_BOOL_ARRAY_TRUE elif isinstance ( idx , NDArrayBase ) and idx . ndim == 0 and idx . dtype != np . bool_ : # This handles ndarray of zero dim. e.g array(1) # while advoid converting zero dim boolean array # float type will be converted to int nonell_key . append ( int ( idx . item ( ) ) ) axis += 1 elif isinstance ( idx , NDArrayBase ) and idx . dtype == np . bool_ : # Necessary size check before using `nonzero` check_boolean_array_dimension ( shape , axis , idx . shape ) # If the whole array is false and npx.set_np() is not set_up # the program will throw infer shape error if not is_np_array ( ) : raise ValueError ( 'Cannot perform boolean indexing in legacy mode. Please activate' ' numpy semantics by calling `npx.set_np()` in the global scope' ' before calling this function.' ) # Add the arrays from the nonzero result to the index nonell_key . extend ( idx . nonzero ( ) ) axis += idx . ndim else : nonell_key . append ( idx ) axis += 1 nonell_key = tuple ( nonell_key ) if ell_idx is None : # This handles the case of "too few" indices, e.g., `nd.zeros((2, 3))[0]`, # where the ellipsis is implicitly after the last entry. ell_idx = len ( nonell_key ) ell_ndim = len ( shape ) + num_none - len ( nonell_key ) expanded_key = ( nonell_key [ : ell_idx ] + ( slice ( None ) , ) * ell_ndim + nonell_key [ ell_idx : ] ) return expanded_key , prepend