def __init__(self, source_iterator: CheckpointableIterator, step_function: Callable[[Any,Any], Tuple[Any,Any]], initial_state: Any = None):
        """
        Args:
            source_iterator: checkpointable iterator to recur over
            step_function: user-supplied function with signature step_function(state, item) -> (new_state, output)
            initial_state: initial state to be passed to the step_function upon first invocation
        """
        if not isinstance(source_iterator, CheckpointableIterator):
            raise ValueError('source_iterator has to be a CheckpointableIterator')
        self._source_iterator = source_iterator               # type: CheckpointableIterator
        self._step_function   = step_function                 # type: Callable[[Any,Any], Tuple[Any,Any]]
        # take deepcopy of initial state so that user cannot change initial state after iterator is created
        self._initial_state   = copy.deepcopy(initial_state)  # type: Any
        self.setstate(None)

        Args:
            source_iterator: checkpointable iterator to recur over
            step_function: user-supplied function with signature step_function(state, item) -> (new_state, output)
            initial_state: initial state to be passed to the step_function upon first invocation
        
kosmos-2/infinibatch/infinibatch/iterators.py