I think this method should work for what you want, although maybe we could make it a bit clearer.
This annotation method
@reg.register("nn.batch_norm", "FTVMExternalCompiler")
def batch_norm(attrs, args, compiler):
if compiler != "dnnl":
return False
return # check with attrs and args
allows you to not just say True of False based on the compiler, but also based on the attrs and args. You could query these to see the tensor shapes and add in a check to say it’s not supported if the input tensors are too large (or other things like if the kernel size isn’t supported).