Repository: dask/dask

Issue:
Per-partition information in dask.dataframe.DataFrame.memory_usage
We currently provide a per-column information in a `pandas.Series` when calling `dask.dataframe.DataFrame.memory_usage` Series. It would also be nice to to have that information on a per-partition-per-column basis. This would allow easier analysis in the case of out-of-memory errors.

I could open a PR that returns this information via a parameter but this would then also change the return-type to a `pandas.DataFrame`. Thus I'm unsure of the right API to implement:

1. `dd.DataFrame.memory_usage(per_partition=True) -> pandas.DataFrame`
2. `dd.DataFrame.memory_usage_detailed() -> pandas.DataFrame`.

Any opinions before I start a PR?

Discussion:
jrbourbeau: Ah good catch, thanks @TomAugspurger! 
jrbourbeau: The Python 3.7 failure is unrelated and appears to sporadically pop up in `master` (xref https://travis-ci.org/dask/dask/jobs/657521668) 
xhochy: I checked out the code locally and ran it over a dataframe with near-equal sized partitions and still see per-column values:

![image](https://user-images.githubusercontent.com/70274/75865682-37f67400-5e04-11ea-812c-2ec95b75bf24.png)

This doesn't match the docstring.
TomAugspurger: Seems like a flaky test with distributed. Opened https://github.com/dask/dask/issues/5972 for that and restarted the build.

Patch:
diff --git a/dask/array/routines.py b/dask/array/routines.py
index 03ea568df..e4668d466 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -239,7 +239,7 @@ def tensordot(lhs, rhs, axes=2):
     if isinstance(axes, Iterable):
         left_axes, right_axes = axes
     else:
-        left_axes = tuple(range(lhs.ndim - 1, lhs.ndim - axes - 1, -1))
+        left_axes = tuple(range(lhs.ndim - axes, lhs.ndim))
         right_axes = tuple(range(0, axes))
 
     if isinstance(left_axes, Integral):
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 0621babe3..e917226e9 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -744,6 +744,29 @@ Dask Name: {name}, {task} tasks"""
 
         return map_overlap(func, self, before, after, *args, **kwargs)
 
+    def memory_usage_per_partition(self, index=True, deep=False):
+        """ Return the memory usage of each partition
+
+        Parameters
+        ----------
+        index : bool, default True
+            Specifies whether to include the memory usage of the index in
+            returned Series.
+        deep : bool, default False
+            If True, introspect the data deeply by interrogating
+            ``object`` dtypes for system-level memory consumption, and include
+            it in the returned values.
+
+        Returns
+        -------
+        Series
+            A Series whose index is the parition number and whose values
+            are the memory usage of each partition in bytes.
+        """
+        return self.map_partitions(
+            total_mem_usage, index=index, deep=deep
+        ).clear_divisions()
+
     @insert_meta_param_description(pad=12)
     def reduction(
         self,
@@ -5654,7 +5677,7 @@ def repartition_size(df, size):
         size = parse_bytes(size)
     size = int(size)
 
-    mem_usages = df.map_partitions(total_mem_usage).compute()
+    mem_usages = df.map_partitions(total_mem_usage, deep=True).compute()
 
     # 1. split each partition that is larger than partition_size
     nsplits = 1 + mem_usages // size
@@ -5675,8 +5698,8 @@ def repartition_size(df, size):
     return _repartition_from_boundaries(df, new_partitions_boundaries, new_name)
 
 
-def total_mem_usage(df):
-    mem_usage = df.memory_usage(deep=True)
+def total_mem_usage(df, index=True, deep=False):
+    mem_usage = df.memory_usage(index=index, deep=deep)
     if is_series_like(mem_usage):
         mem_usage = mem_usage.sum()
     return mem_usage
diff --git a/docs/source/dataframe-api.rst b/docs/source/dataframe-api.rst
index be6786a33..74a389775 100644
--- a/docs/source/dataframe-api.rst
+++ b/docs/source/dataframe-api.rst
@@ -49,6 +49,8 @@ Dataframe
     DataFrame.mask
     DataFrame.max
     DataFrame.mean
+    DataFrame.memory_usage
+    DataFrame.memory_usage_per_partition
     DataFrame.merge
     DataFrame.min
     DataFrame.mod
@@ -162,6 +164,7 @@ Series
    Series.max
    Series.mean
    Series.memory_usage
+   Series.memory_usage_per_partition
    Series.min
    Series.mod
    Series.mul


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index 952c082f7..f192caff6 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -304,6 +304,31 @@ def test_tensordot_2(axes):
     assert_eq(da.tensordot(y, y, axes=axes), np.tensordot(x, x, axes=axes))
 
 
+@pytest.mark.parametrize("chunks", ["auto", (4, 6), (2, 3), (4, 3), (2, 6)])
+def test_tensordot_double_contraction_neq2(chunks):
+    # Regression test for https://github.com/dask/dask/issues/5472
+    x = np.arange(24).reshape(4, 6)
+    y = da.from_array(x, chunks=chunks)
+    assert_eq(da.tensordot(y, y, axes=2), np.tensordot(x, x, axes=2))
+
+
+def test_tensordot_double_contraction_ngt2():
+    # Regression test for https://github.com/dask/dask/issues/5472
+    x = np.arange(60.0).reshape(3, 4, 5)
+    y = np.arange(60.0).reshape(4, 5, 3)
+    u = da.from_array(x)
+    v = da.from_array(y)
+
+    assert_eq(da.tensordot(u, v, axes=2), np.tensordot(x, y, axes=2))
+
+    x = np.arange(60.0).reshape(3, 4, 5)
+    y = np.arange(60.0).reshape(4, 5, 3)
+    u = da.from_array(x, chunks=3)
+    v = da.from_array(y)
+
+    assert_eq(da.tensordot(u, v, axes=2), np.tensordot(x, y, axes=2))
+
+
 def test_tensordot_more_than_26_dims():
     ndim = 27
     x = np.broadcast_to(1, [2] * ndim)
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index aa11ce249..f2279f51f 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -1802,7 +1802,7 @@ def test_repartition_partition_size(use_index, n, partition_size, transform):
     a = dd.from_pandas(df, npartitions=n, sort=use_index)
     b = a.repartition(partition_size=partition_size)
     assert_eq(a, b, check_divisions=False)
-    assert np.alltrue(b.map_partitions(total_mem_usage).compute() <= 1024)
+    assert np.alltrue(b.map_partitions(total_mem_usage, deep=True).compute() <= 1024)
     parts = dask.get(b.dask, b.__dask_keys__())
     assert all(map(len, parts))
 
@@ -3614,6 +3614,34 @@ def test_memory_usage(index, deep):
     )
 
 
+@pytest.mark.parametrize("index", [True, False])
+@pytest.mark.parametrize("deep", [True, False])
+def test_memory_usage_per_partition(index, deep):
+    df = pd.DataFrame(
+        {
+            "x": [1, 2, 3, 4, 5],
+            "y": [1.0, 2.0, 3.0, 4.0, 5.0],
+            "z": ["a", "b", "c", "d", "e"],
+        }
+    )
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    # DataFrame.memory_usage_per_partition
+    expected = pd.Series(
+        part.compute().memory_usage(index=index, deep=deep).sum()
+        for part in ddf.partitions
+    )
+    result = ddf.memory_usage_per_partition(index=index, deep=deep)
+    assert_eq(expected, result)
+
+    # Series.memory_usage_per_partition
+    expected = pd.Series(
+        part.x.compute().memory_usage(index=index, deep=deep) for part in ddf.partitions
+    )
+    result = ddf.x.memory_usage_per_partition(index=index, deep=deep)
+    assert_eq(expected, result)
+
+
 @pytest.mark.parametrize(
     "reduction",
     [