Refactor ops to allow non-funsor parameters

[fritzo]

Addresses #489 partly pair coded with @eb8680

This major op refactoring aims to unblock @ordabayevy in #477 and #482 by allowing ops to have arbitrary *arg,**kwarg parameters. Changes include:

• nullop is renamed to null for consistency (NullOp remains unchanged).
• Dispatch tables are moved from op instances to op subclasses.
• @make_op has moved to a classmethod @MyOpClass.make e.g. @BinaryOp.make
• Each concrete op subclass has a fixed .arity attribute.
• Op patterns dispatch on only the first [:arity]-many args (and patterns are validated to ensure this).
• Ops are new weakly cached via WeakValueDict.
• .subclass_register() handlers have a different inferface, now including cls.
• There is now a single shared WrappedTransformOp class, and its instances each have an op.
• ops.WrappedTransformOp(fn) must now be written as a kwarg ops.WrappedTransformOp(fn=fn)
• is_numeric_array() is demoted from an Op to a @singledispatch.
• ops.stack interface has changed to ops.stack(parts, dim) as in NumPy and PyTorch.
• funsor.tensor.stack is replaced by ops.stack.
• ops.einsum interface has changed to ops.einsum(operands, equation) where operands is a list or tuple.
• Finitary funsor interface has changed from variously-many args to a single tuple arg.

Tested

• refactoring is covered by existing tests

[fritzo]

@eb8680 finally got the tests to pass 😓

[fehiepsi]

For some reason, I am getting

E       AttributeError: 'DirichletMultinomial' object has no attribute 'data'

in NumPyro test/test_handlers.py::test_collapse_beta_binomial. The instance DirichletMultinomial has the value

DirichletMultinomial(concentration=[1.5 0.5], total_count=10.0, value=Finitary(stack, (Tensor(7.0, OrderedDict(), 'real'), Number(3.0))))

[fritzo]

@eb8680, @fehiepsi's issue with DirichletMultinomial sounds related to the lazy test we disabled. Do you think we should try to get that working again?

[fehiepsi]

Yeah, enabling it raises two errors

• NotImplementedError: inhomogeneous total_count is not supported (only with torch backend)
• NotImplementedError: TODO support converting Indep(Transform)

but it seems to me refactoring of stack is more relevant.