Honestly, it's a little suspicious whether we should even allow this: should being diagonal perhaps imply that the input and output structures are identical?
Right now I'm choosing to allow this because it's pretty subtle, and users can apply their own diagonal tags, so if nothign else it's an easy mistake to be tolerant of.
In particular, we were making this mistake by treating scalar operators as diagonal even when they had different structures. This was causing a downstream issue in Diffrax.
Honestly, it's a little suspicious whether we should even allow this: should being diagonal perhaps imply that the input and output structures are identical?
Right now I'm choosing to allow this because it's pretty subtle, and users can apply their own diagonal tags, so if nothign else it's an easy mistake to be tolerant of.
In particular, we were making this mistake by treating scalar operators as diagonal even when they had different structures. This was causing a downstream issue in Diffrax.