Add tests for all callbacks types

[scarlehoff]

Right now we are compatible with some of the callbacks types of TF but we should test (and add CI tests) for all callbacks types to ensure they continue to work as expected.