Add tests and opt one more case.

Review me code please before accepting it.

If you merge it, I'll merge the PR in Theano.