It is a little annoying that you have to expand the dims to write (as shown in readme). Would be nice to just write as is.