Support output_spec for custom batch functions.
#1292
Merged
output_spec for custom batch functions.
#1292