Improve error handling in callbacks when users use incorrect type annotations

[erick-xanadu]

This may lead to garbage memory being used or overwriting memory. I think this can be caught doing static type checking in MLIR. Probably during the lowering of backprop op.