doctests with use_beam=False does not partition the output expression

We call expression.evaluate_at, rather than session.evaluate. The former doesn't go through PartitioningSession for the output expression, so that expression is never partitioned. This is covering up some argument configurations that we shouldn't allow.