Consider forking ubersum/tensor contractions?

srush · October 28, 2018, 4:55pm

I am really liking the work going on in discrete models and the extensions to einsum. It’s really a beautiful way to think about discrete graphical models, particularly the multi-expression right-hand side syntax.

Any chance you would consider forking it out to a separate project as with distributions? We have a lot of uses for it independent of Pyro, where I am hesitant to want a full Pyro dependency.

Thanks,
Sasha

fritzo · October 28, 2018, 5:15pm

Hi @srush, glad you find ubersum() useful! It was originally Eli’s idea to think of discrete graphical models this way.

Yes we plan to push the implementation upstream to opt_einsum once it is correct (I’m working on fixing the last known bug today). We’ve already pushed the sharing logic upstream, so if you just need multiple outputs you can write a simple wrapper. The remaining work beyond that simple wrapper is to query batches of local variables that communicate through a global variable.

srush · October 28, 2018, 6:22pm

Sweet! That is ideal. I’m gonna build a bunch of neural CRFs.

One more small thing. This line has me a bit worried:

github.com

pyro-ppl/pyro/blob/dev/pyro/ops/einsum/torch_log.py#L5


      
          # Copyright (c) 2017-2019 Uber Technologies, Inc.
          # SPDX-License-Identifier: Apache-2.0
          
          
import torch
          
          
from pyro.ops.einsum.util import Tensordot
          from pyro.ops.special import safe_log
          
          

          
def transpose(a, axes):
              return a.permute(axes)
          
          

          
def einsum(equation, *operands):
              """

EINSUM_SYMBOLS_BASE = ‘abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ’

Just want to be sure that this will all support arbitrary size CUDA variables. Seems like opt-einsum does that:

https://optimized-einsum.readthedocs.io/en/latest/ex_large_expr_with_greedy.html
get_symbol(i)

srush · October 28, 2018, 7:13pm

Oops sorry, I misunderstood how opt-einsum works. Now I see that it has its own way of symbol mapping internally. Very cool.