David Choi

Would it be possible to make [this check](https://github.com/google/jax/blob/837f0bbf6f667bacb85243e84d9a75e6b6eef1c7/jax/_src/custom_derivatives.py#L777) introduced [here](https://github.com/google/jax/commit/1326c74db7f25b7224412fbe6f69f43d17a90e6a) optional? We were previously putting things with deliberately different shapes into this for logging, as an easy way of extracting...