add numpy einsum backend #314
Conversation
|
This should be unblocked now, right? |
funsor/einsum/util.py
Outdated
| x = ops.permute(x, tuple(old_dims.index(dim) for dim in dims if dim in old_dims)) | ||
| x = x.reshape(tuple(sizes[dim] if dim in old_dims else 1 for dim in dims)) | ||
| x = ops.expand(x, shape) | ||
| # workaround: ndarray does not allow setting attribute "_pyro_dims" |
There was a problem hiding this comment.
I think this is an important point: we can't set new attributes to ndarray, DeviceArray directly.
|
@eb8680 has additional comments, please don't merge yet |
funsor/einsum/numpy_map.py
Outdated
| This assumes all operands have a ``._pyro_dims`` attribute set. | ||
| """ | ||
| # TODO: the implementation skips "backward" logic | ||
| equation = rename_equation(equation, *operands) |
There was a problem hiding this comment.
Is this actually necessary now? What happens if you remove this line?
There was a problem hiding this comment.
Sorry I missed your comments. I'll address those soon.
There was a problem hiding this comment.
Yes, it is necessary if we keep _pyro_dims attribute. One of the test_einsum fail: test_einsum[funsor.einsum.numpy_map-ab,bc,cd->da].
funsor/einsum/numpy_map.py
Outdated
|
|
||
| contract_dims = ''.join(sorted(set().union(*(x._pyro_dims for x in operands)) - set(output))) | ||
| dims = output + contract_dims | ||
| result = reduce(operator.add, broadcast_all(*operands, dims=dims)) |
There was a problem hiding this comment.
Is this the only place where _pyro_dims are actually used and not just passed along?
There was a problem hiding this comment.
This and rename_equation, broadcast_all utility require pyro_dims. I am trying to remove pyro_dims. If things pass, then I think we can simplify much code here.
|
@eb8680 To remove |
|
@fehiepsi sure, that sounds easy enough |
|
@eb8680 It seems that it is actually easy. Thanks for guiding! |
This PR ports pyro einsum MAP backend to funsor. Here are changes in this PR:
funsor.einsum.numpy_log,funsor.einsum.numpy_mapbackend-agnostic.set_backend(get_backend())infunsor.__init__file due to a circular import issue (I'll revisit this issue later).There are some issues with numpy backend:
._pyro_dims(and._backwardstuff) to ndarray/DeviceArray.memoizetests fail withnumpy.genericdata. As a workaround, I castnumpy.generictonumpy.ndarrayat Tensor construction. Because this issue does not happen with JAX backend, I didn't attempt to resolve it.