batch_dims seems not to be respected for 0.2.0?

[arbennett]

What happened?

I believe this is somewhat related to #121 where up front computation is slowed down due to ignoring the batch_dims keyword argument, meaning that batches are essentially duplicated, and very large to compute. This is reproducible in a google colab notebook just by changing the installed version between 0.1.0 and 0.2.0.

Minimal Complete Verifiable Example

import xbatcher
import xarray as xr
import numpy as np
shape = (30, 4, 2)
dims = ('z', 'y', 'x')

ds = xr.Dataset()
ds['v1'] = xr.DataArray(data=np.random.random(shape), dims=dims )
ds['v2'] = xr.DataArray(data=np.random.random(shape), dims=dims )
bgen1 =  xbatcher.BatchGenerator(
    ds=ds, input_dims={'z': 10}
)
bgen2 =  xbatcher.BatchGenerator(
    ds=ds, input_dims={'z': 10}, batch_dims={'y': 2, 'x': 1},
)

print('No batch_dims:   ', next(iter(bgen1)).dims)
print('With batch_dims: ', next(iter(bgen2)).dims)

Environment

On version 0.1.0 I get the output:

No batch_dims:    Frozen({'z': 10, 'sample': 8})
With batch_dims:  Frozen({'z': 10, 'sample': 2})

On version 0.2.0 I get the output:

No batch_dims:    Frozen({'z': 10, 'sample': 8})
With batch_dims:  Frozen({'z': 10, 'sample': 8})

[maxrjones]

Thanks for the bug report and MVCE @arbennett.

Yes, this is connected to #121, with an explanation of the problem in #121 (comment). In summary, batch_dims in v0.2.0 erroneously only controls the number of batches but not their size. I'm aiming to open a proposed fix today.

[maxrjones]

As a stopgap measure, if needed you could install the last pre-regression commit using:

python -m pip install git+https://github.com/xarray-contrib/xbatcher.git@4d8e2c84a2d405e237f60f1df5286dd766e06ff0