update moments for normalize bug (NVIDIA-Merlin#444)

rnyak · web-flow · commit 94e9402fae5b · 2020-11-20T09:51:29.000-08:00
diff --git a/nvtabular/ops/moments.py b/nvtabular/ops/moments.py
@@ -125,7 +125,7 @@ def _chunkwise_moments(df):
         df2[col] = df[col].astype("float64").pow(2)
     vals = {
         "df-count": df.count().to_frame().transpose(),
-        "df-sum": df.sum().to_frame().transpose(),
+        "df-sum": df.sum().astype("float64").to_frame().transpose(),
         "df2-sum": df2.sum().to_frame().transpose(),
     }
     # NOTE: Perhaps we should convert to pandas here
diff --git a/tests/unit/test_ops.py b/tests/unit/test_ops.py
@@ -436,6 +436,38 @@ def test_normalize(tmpdir, df, dataset, gpu_memory_frac, engine, op_columns):
     assert new_gdf["x"].equals(df["x"])
 
 
+@pytest.mark.parametrize("gpu_memory_frac", [0.1])
+@pytest.mark.parametrize("engine", ["parquet"])
+@pytest.mark.parametrize("op_columns", [["x"], None])
+def test_normalize_upcastfloat64(tmpdir, dataset, gpu_memory_frac, engine, op_columns):
+    df = cudf.DataFrame(
+        {"x": [1.9e10, 2.3e16, 3.4e18, 1.6e19], "label": [1, 0, 1, 0]}, dtype="float32"
+    )
+
+    cat_names = []
+    cont_names = ["x"]
+    label_name = ["label"]
+
+    config = nvt.workflow.get_new_config()
+    config["PP"]["continuous"] = [ops.Moments(columns=op_columns)]
+
+    processor = nvtabular.Workflow(
+        cat_names=cat_names, cont_names=cont_names, label_name=label_name, config=config
+    )
+
+    processor.update_stats(dataset)
+
+    op = ops.Normalize()
+
+    columns_ctx = {}
+    columns_ctx["continuous"] = {}
+    columns_ctx["continuous"]["base"] = op_columns or cont_names
+
+    new_gdf = op.apply_op(df, columns_ctx, "continuous", stats_context=processor.stats)
+    df["x"] = (df["x"] - processor.stats["means"]["x"]) / processor.stats["stds"]["x"]
+    assert new_gdf["x"].equals(df["x"])
+
+
 @pytest.mark.parametrize("gpu_memory_frac", [0.01, 0.1])
 @pytest.mark.parametrize("engine", ["parquet", "csv", "csv-no-header"])
 @pytest.mark.parametrize("op_columns", [["x"], None])

Original file line number	Diff line number	Diff line change
`@@ -125,7 +125,7 @@ def _chunkwise_moments(df):`
`125`	`125`	`df2[col] = df[col].astype("float64").pow(2)`
`126`	`126`	`vals = {`
`127`	`127`	`"df-count": df.count().to_frame().transpose(),`
`128`		`- "df-sum": df.sum().to_frame().transpose(),`
	`128`	`+ "df-sum": df.sum().astype("float64").to_frame().transpose(),`
`129`	`129`	`"df2-sum": df2.sum().to_frame().transpose(),`
`130`	`130`	`}`
`131`	`131`	`# NOTE: Perhaps we should convert to pandas here`