docs: 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)