Skip to content

Commit

Permalink
eager categorical constant pattern (#575)
Browse files Browse the repository at this point in the history
  • Loading branch information
ordabayevy authored Oct 18, 2021
1 parent c94a9bd commit 8aec9f4
Showing 1 changed file with 4 additions and 0 deletions.
4 changes: 4 additions & 0 deletions funsor/torch/distributions.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@

import funsor.ops as ops
from funsor.cnf import Contraction
from funsor.constant import Constant
from funsor.distribution import ( # noqa: F401
FUNSOR_DIST_NAMES,
Bernoulli,
Expand Down Expand Up @@ -362,6 +363,9 @@ def deltadist_to_funsor(pyro_dist, output=None, dim_to_name=None):
eager.register(Multinomial, Tensor, Tensor, Tensor)(eager_multinomial) # noqa: F821)
eager.register(Categorical, Funsor, Tensor)(eager_categorical_funsor) # noqa: F821)
eager.register(Categorical, Tensor, Variable)(eager_categorical_tensor) # noqa: F821)
eager.register(Categorical, Constant[Tuple, Tensor], Variable)(
eager_categorical_tensor
) # noqa: F821)
eager.register(Delta, Tensor, Tensor, Tensor)(eager_delta_tensor) # noqa: F821
eager.register(Delta, Funsor, Funsor, Variable)(
eager_delta_funsor_variable
Expand Down

0 comments on commit 8aec9f4

Please sign in to comment.