* fixed merge conflicts

* fixed failing tests and added rewrite for pt.diag

* minor changes; added test to not apply rewrite

* added test for batched case and more cases of not applying rewrite

* minor changes