Add ReductionOp, ops.mean, ops.std, ops.var

[ordabayevy]

No description provided.

[fritzo]

Could you register implementations in funsor/jax/ops.py the following test to test/test_terms.py:

diff --git a/test/test_terms.py b/test/test_terms.py
index e495e3e..88ed47a 100644
--- a/test/test_terms.py
+++ b/test/test_terms.py
@@ -287,6 +287,29 @@ def test_unary(symbol, data):
     check_funsor(actual, {}, Array[dtype, ()], expected_data)


+@pytest.mark.parametrize("event_shape", [(4,), (3, 2)], ids=str)
+@pytest.mark.parametrize(
+    "name",
+    [
+        "all",
+        "any",
+        "logsumexp",
+        "max",
+        "mean",
+        "min",
+        "prod",
+        "std",
+        "sum",
+        "var",
+    ],
+)
+def test_reduce_event(name, event_shape):
+    dtype = 2 if name in ("any", "all") else "real"
+    x = random_tensor(OrderedDict(i=Bint[5]), output=Array[dtype, event_shape])
+    actual = getattr(x, name)()
+    check_funsor(actual, x.inputs, Array[dtype, ()])
+
+
 BINARY_OPS = [
     "+",
     "-",

[ordabayevy]

Is there a need to allow ddof option for ops.std and ops.var? If yes, how is it passed to Unary? Similar to shape option in .reshape?

    def reshape(self, shape):
        return Unary(ops.ReshapeOp(shape), self)

[fritzo]

Is there a need to allow ddof option

I'm not sure what the ddof option is, but it would be reasonable to create, in addition to these default mean,var,std ops, more flexible op classes MeanOp, VarOp, and StdOp to take non-default arguments (axis,keepdims) (in numpy parlance) or (dim,keepdim) (in pytorch parlance). You're correct that these could follow ReshapeOp or GetitemOp by using CachedOpMeta. However I think those ops would best be separate and complementary to the default ops, and we would also want to provide them for sum etc.:

class Funsor(...):
    ...
    def mean(self, axis=None, keepdims=False):
        if axis is None and keepdims is False:
            op = ops.mean  # default version from this PR
        else:
            op = ops.MeanOp(axis, keepdims)  # fancy version from a future PR
        return Unary(op, self)

Or if you refactored such that ops.mean = MeanOp(), you could simplify to

class Funsor(...):
    ...
    def mean(self, axis=None, keepdims=False):
        return Unary(ops.MeanOp(axis, keepdims), self)

[fritzo]

Looks great! I'm still uncertain of argmin, argmax...

[ordabayevy]

I'm still uncertain of argmin, argmax...

There were some rough edges and argmin and argmax are definitely one of them. Probably should be in a separate ArgreductionOp as you suggested.

[fritzo]

🎉