torch reduce fns: no default dim or keepdim value

Author: jcmgray

autoray/autoray.py

@@ -1312,21 +1312,24 @@ def translated_function(*args, **kwargs):
         new_kwargs = {}
         translation = translator.copy()
 
-        # convert args
+        # convert args, pairing them off with kwargs
         for arg_value in args:
             new_arg_name = translation.popitem(last=False)[1][0]
             new_kwargs[new_arg_name] = arg_value
 
-        # convert kwargs -  but only those in the translation
+        # convert kwargs - but only those in the translation
         for key, value in kwargs.items():
             try:
                 new_kwargs[translation.pop(key)[0]] = value
             except KeyError:
                 new_kwargs[key] = value
 
         # set remaining default kwargs
-        for key, value in translation.items():
-            new_kwargs[value[0]] = value[1]
+        for opt in translation.values():
+            if len(opt) == 2:
+                # backend_name, default_value
+                new_kwargs[opt[0]] = opt[1]
+            # else, no default value -> don't inject
 
         return fn(**new_kwargs)
 
@@ -2146,12 +2149,13 @@ def numpy_like(x, axis=None):
 _CUSTOM_WRAPPERS["torch", "sort"] = torch_sort_wrap
 _CUSTOM_WRAPPERS["torch", "flip"] = torch_flip_wrap
 _torch_reduce_translation = [
-    ("a", ("input",)), ("axis", ("dim", None)), ("keepdims", ("keepdim", False))
+    ("a", ("input",)),
+    ("axis", ("dim",)),
+    ("keepdims", ("keepdim",)),
 ]
-_CUSTOM_WRAPPERS["torch", "sum"] = make_translator(_torch_reduce_translation)
-_CUSTOM_WRAPPERS["torch", "max"] = make_translator(_torch_reduce_translation)
-_CUSTOM_WRAPPERS["torch", "min"] = make_translator(_torch_reduce_translation)
-_CUSTOM_WRAPPERS["torch", "prod"] = make_translator(_torch_reduce_translation)
+for f in ("sum", "max", "min", "prod", "mean", "median", "std", "var"):
+    # TODO: search "keepdim" in torch docs to find more
+    _CUSTOM_WRAPPERS["torch", f] = make_translator(_torch_reduce_translation)
 
 # for older versions of torch, can provide some alternative implementations
 _MODULE_ALIASES["torch[alt]"] = "torch"

tests/test_autoray.py

@@ -149,6 +149,42 @@ def test_binary_functions(f, args, xfail_backends, backend):
     assert ar.do("allclose", yt, yn)
 
 
+@pytest.mark.parametrize(
+    "f",
+    [
+        "sum",
+        "prod",
+        "max",
+        "min",
+        "mean",
+    ],
+)
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {},
+        {"axis": 1},
+        {"axis": 1, "keepdims": True},
+        {"axis": (0, 2)},
+    ],
+)
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_reduce_functions(f, kwargs, backend):
+    if (
+        backend == "torch"
+        and f == "prod"
+        and isinstance(kwargs.get("axis"), tuple)
+    ):
+        pytest.xfail("Pytorch doesn't support prod with tuple axis.")
+
+    x = ar.do("random.normal", size=(2, 3, 4), like="numpy")
+    y = ar.do(f, x, **kwargs)
+    xb = ar.do("asarray", x, like=backend)
+    yb = ar.do(f, xb, **kwargs)
+    yt = ar.do("to_numpy", yb)
+    assert ar.do("allclose", yt, y)
+
+
 @pytest.mark.parametrize("backend", BACKENDS)
 @pytest.mark.parametrize("fn", ["sqrt", "exp", "sum"])
 def test_basic(backend, fn):