Support use under PyTorch jit

[fritzo]

Description

This fixes path optimizers to work when run under the PyTorch jit. The problem was that when running under the jit, x.shape is a tensor, and x.shape[pos] is also a tensor. These tensors can't be operated on as easily as ints. this PR merely casts to an int. Note this will still raise a warning in the PyTorch jit, but it is easy to filter out the warning; typically we want to optimize the path once, then replay compiled code multiple times assume tensor shapes are roughly similar.

Tested

• ran locally, showing that this fixes some jit usage in Pyro

Status

• Ready to go

[codecov-io]

Codecov Report

Merging #77 into master will not change coverage.
The diff coverage is 100%.

[dgasmith]

Interesting, thanks for the fix.

[dgasmith]

Will you need a 2.3.x release for this?

[fritzo]

Will you need a 2.3.x release for this?

A 2.3.x release would be great 👍 We currently patch opt_einsum to work around this issue, and we can remove the gross patch as soon as opt_einsum is released. Thanks!

[dgasmith]

Ok, 2.3.2 out with this patch in on PyPI. Should on conda-forge in a few hours.

[fritzo]

Thanks, @dgasmith !