Parallelizing 3D numpy array calculation using dask.array.core.map_blocks

I have a 3D numpy array (dimensions: depth, latitude, longtitude) and I am trying to do some parallelized calculation using the data along the depth axis at each lat-lon point and so far I have been unsuccessful. I've looked at the documentation for dask.array.core.map_blocks but it has not been helpful. Here is what I am doing:

N2_dask = da.from_array(N2_naned, chunks=(49, 32, (12, 12, 12, 12, 12, 12)))
zN2_dask = da.from_array(-zN2_agg[t], chunks=(49, 32, (12, 12, 12, 12, 12, 12)))
lat_dask = da.from_array(lat_agg, chunks=(32))
lon_dask = da.from_array(lon_agg, chunks=((12, 12, 12, 12, 12, 12)))

for j in range(len(lat_dask)):
    for i in range(len(lon_dask)):
        f = da.core.map_blocks(baroclinic.neutral_modes_from_N2_profile(
            zN2_dask[:, j, i], N2_dask[:, j, i], gsw.earth.f(lat_dask[j, i]), **kwargs))
        zphi, Rd, vd = f.compute()

where baroclinic.neutral_modes_from_N2_profile is my function. I get an error as the follows:

AssertionErrorTraceback (most recent call last)
<ipython-input-61-e58a7b54c470> in <module>()
      6         print zN2_dask[:, j, i]
      7         f = da.core.map_blocks(baroclinic.neutral_modes_from_N2_profile(
----> 8                 zN2_dask[:, j, i], N2_dask[:, j, i], gsw.earth.f(lat_dask[j, i]), **kwargs))
      9         zphi, Rd, vd = f.compute()

/home/takaya/.conda/envs/oceanmodes/lib/python2.7/site-packages/dask/array/core.pyc in __getitem__(self, index)
   1023             return self
   1024 
-> 1025         dsk, chunks = slice_array(out, self.name, self.chunks, index)
   1026 
   1027         return Array(merge(self.dask, dsk), out, chunks, dtype=self._dtype)

/home/takaya/.conda/envs/oceanmodes/lib/python2.7/site-packages/dask/array/slicing.pyc in slice_array(out_name, in_name, blockdims, index)
    134 
    135     # Pass down to next function
--> 136     dsk_out, bd_out = slice_with_newaxes(out_name, in_name, blockdims, index)
    137 
    138     bd_out = tuple(map(tuple, bd_out))

/home/takaya/.conda/envs/oceanmodes/lib/python2.7/site-packages/dask/array/slicing.pyc in slice_with_newaxes(out_name, in_name, blockdims, index)
    152 
    153     # Pass down and do work
--> 154     dsk, blockdims2 = slice_wrap_lists(out_name, in_name, blockdims, index2)
    155 
    156     # Insert ",0" into the key:  ('x', 2, 3) -> ('x', 0, 2, 0, 3)

/home/takaya/.conda/envs/oceanmodes/lib/python2.7/site-packages/dask/array/slicing.pyc in slice_wrap_lists(out_name, in_name, blockdims, index)
    183     shape = tuple(map(sum, blockdims))
    184     assert all(isinstance(i, (slice, list, int, long)) for i in index)
--> 185     assert len(blockdims) == len(index)
    186     for bd, i in zip(blockdims, index):
    187         check_index(i, sum(bd))

AssertionError: 

Can anyone tell me why this gives me an AssertionError? Thank you in advance!