code: # Copyright 2019 The FastEstimator Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  from typing import Any, Dict, Iterable, List, Union
  
  import numpy as np
  
  from fastestimator.op.numpyop.numpyop import NumpyOp
  from fastestimator.util.traceability_util import traceable
  
  
  @traceable()
  class ExpandDims(NumpyOp):
      """Expand the dimension of inputs by inserting a new axis to the specified axis position.
  
      Args:
          inputs: Key(s) of inputs to be modified.
          outputs: Key(s) into which to write the modified inputs.
          mode: What mode(s) to execute this Op in. For example, "train", "eval", "test", or "infer". To execute
              regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument
              like "!infer" or "!train".
          axis: The axis to expand.
      """
      def __init__(self,
                   inputs: Union[str, Iterable[str]],
                   outputs: Union[str, Iterable[str]],
                   mode: Union[None, str, Iterable[str]] = None,
                   axis: int = -1):
          super().__init__(inputs=inputs, outputs=outputs, mode=mode)
          self.axis = axis
          self.in_list, self.out_list = True, True
  
      def forward(self, data: List[np.ndarray], state: Dict[str, Any]) -> List[np.ndarray]:
          return [np.expand_dims(elem, self.axis) for elem in data]
  
apis: numpy.expand_dims