Skip to content

Commit 6cd22dd

Browse files
author
Brendan Cooley
committed
fix(distributions.conjugate): add total_count_max to pytree_aux_fields
1 parent 675c2b3 commit 6cd22dd

File tree

5 files changed

+4
-1
lines changed

5 files changed

+4
-1
lines changed

ar2.png

40.8 KB
Loading

graph

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
digraph {
2+
x [label=x fillcolor=white shape=ellipse style=filled]
3+
}

numpyro/distributions/conjugate.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ class DirichletMultinomial(Distribution):
115115
"total_count": constraints.nonnegative_integer,
116116
}
117117
pytree_data_fields = ("concentration", "_dirichlet")
118-
pytree_aux_fields = ("total_count",)
118+
pytree_aux_fields = ("total_count", "total_count_max")
119119

120120
def __init__(
121121
self, concentration, total_count=1, *, total_count_max=None, validate_args=None

ssbvm_mixture.png

34.3 KB
Loading

wordclouds.png

626 KB
Loading

0 commit comments

Comments
 (0)