These changes "flatten" the nested conditions that lead to the replacement
logic.

They also clarify the tests that are labeled as being for
`local_add_specialize`, which actually aren't.  Also, the unrelated tests
implicitly relied on the canonicalizations built into
`get_scalar_constant_value`; now, the actual canonicalizations are required in
anticipation of `get_scalar_constant_value`'s replacement.

* Convert tests to pytest `parametrized` tests
* Use `pytest.raises`
* Remove unused timing code and methods

Closes #633

JAX can check for correct axis itself