Support ops.sum(data, dim=None, keepdims=False)

[fritzo]

Addresses #489
pair coded with @eb8680 @ordabayevy @fehiepsi

This demonstrates the new parametrized op syntax from #491 . The recipe is:

1. add new *args, *kwargs to your op in funsor.ops.array
2. add backend support in funsor.torch.ops and funsor.jax.ops
3. implement find_domain(op, ...) in funsor.domains
4. work around batch dimensions in funsor.tensor
5. add some tests

Tested

• unit test for various op arg,kwarg syntax
• unit test for batched tensors

[fritzo]

@ordabayevy I hope this can serve as a template for your #482 . Feel free to refactor after this PR merges, in case you'd like to reuse the logic e.g. in find_domain or eager.register. We may even be able to create a subclass ReductionOp(UnaryOp), replace to use

- @UnaryOp.make
+ @ReductionOp.make
  def sum(...):
      ...

and register ops.sum, ops.mean, etc. in a single pattern.

[eb8680]

This logic converting dimensions into batch-aware dimensions seems useful and general-purpose enough that some version of it should maybe live in an Op method or something - we don't want to have to write this from scratch in each new op with nontrivial shape semantics.

[fritzo]

I agree this is useful, and I'm hoping @ordabayevy will generalize this in #482 once we have more than one use case. This file seems like the right place for that general logic since funsor.ops should be agnostic to inputs and domains etc.