Had the same idea; opened a "feature request" here: https://github.com/m4nuC/aifiddle-issues/issues/6

Not sure how easy it will be to support all the types of layers on all platforms, but I think if target is Keras/PyTorch already going to be kick ass...