Reset the state to the given initial state.

        Args:
            reset_state: If this is an int, then this is the state to reset
                the stepper to, expressed as an integer of the computational
                basis. Integer to bitwise indices is little endian. Otherwise
                if this is a np.ndarray this must be the correct size, be
                normalized (L2 norm of 1), and have dtype of np.complex64.

        Raises:
            ValueError if the state is incorrectly sized or not of the correct
            dtype.

def reset_state(self, reset_state):
        """Reset the state to the given initial state.

        Args:
            reset_state: If this is an int, then this is the state to reset
                the stepper to, expressed as an integer of the computational
                basis. Integer to bitwise indices is little endian. Otherwise
                if this is a np.ndarray this must be the correct size, be
                normalized (L2 norm of 1), and have dtype of np.complex64.

        Raises:
            ValueError if the state is incorrectly sized or not of the correct
            dtype.
        """
        # If the pool has been closed, recreate to calculate state.'
        if isinstance(reset_state, int):
            self._pool.map(_reset_state,
                           self._shard_num_args({'reset_state': reset_state}))
        elif isinstance(reset_state, np.ndarray):
            sim.validate_normalized_state(reset_state, self._num_qubits)
            args = []
            for kwargs in self._shard_num_args():
                shard_num = kwargs['shard_num']
                shard_size = 1 << kwargs['num_shard_qubits']
                start = shard_num * shard_size
                end = start + shard_size
                kwargs['reset_state'] = reset_state[start:end]
                args.append(kwargs)
            self._pool.map(_reset_state, args)