Callback helper class for function callbacks during forward and backward class.
This is used from Python frontend.