flax

support batch dimension in image summaries

#1102

Merged