What is the type hint of variables bound inside (e.g., ax in Softmax)

[ordabayevy]

What should be the type hint for ax in this Softmax function?

1. Funsor works but then then ax cannot be a str
2. Value[str] can only be a str and not Variable
3. Bound works both as str and Variable but ax is not actually bound

@make_funsor
def Softmax(
    x: Funsor,
    ax: ? # Funsor or Bound or Value[str]
) -> Fresh[lambda x: x]:
    return x.exp() / x.exp().reduce(ops.add, ax)

[fritzo]

You'll need ax: Bound and ax2: Fresh[lambda ax: ax]. As we discussed yesterday, @make_funsor currently does not support binding-and-return of a dimension with the same name, so for now you'll need to rename the dimension. I believe @eb8680 is planning to address that as part of alpha-renaming work.

@make_funsor
def Softmax(
    x: Funsor,
    ax: Bound,
    ax2: Fresh[lambda ax: ax],
) -> Fresh[lambda x: x]:
    x = x(**{ax: ax2})  # rename to work around alpha renaming limitations
    y = x - x.reduce(ops.logaddexp, ax2)
    return y.exp()