fix type inference for iteration with stagedfunction