def assert_chunks_equal(a, b):
    """
    Assert that chunksizes along chunked dimensions are equal.

    Parameters
    ----------
    a : xarray.Dataset or xarray.DataArray
        The first object to compare.
    b : xarray.Dataset or xarray.DataArray
        The second object to compare.
    """

    if isinstance(a, DataArray) != isinstance(b, DataArray):
        raise TypeError("a and b have mismatched types")

    left = a.unify_chunks()
    right = b.unify_chunks()
    assert left.chunks == right.chunks

    Assert that chunksizes along chunked dimensions are equal.

    Parameters
    ----------
    a : xarray.Dataset or xarray.DataArray
        The first object to compare.
    b : xarray.Dataset or xarray.DataArray
        The second object to compare.
    
xarray/testing/assertions.py