experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Adds double support for StatelessMultinomial (which is really a categorical).

Browse Source 
Relevant to https://github.com/tensorflow/probability/issues/1127

PiperOrigin-RevId: 336903719
Change-Id: Id46908a9e57874387ef77208a21eb6ce34e6a9b1
This commit is contained in:
Brian Patton 2020-10-13 10:22:55 -07:00 committed by TensorFlower Gardener
parent 4c6e31c085
commit dd21d56ae8
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/compiler/tf2xla/kernels/categorical_op.cc
View File

@ -186,7 +186,7 @@ class StatelessCategoricalOp : public CategoricalOp {
REGISTER_XLA_OP(Name("StatelessMultinomial")
.CompileTimeConstantInput("num_samples")
.TypeConstraint("T", {DT_FLOAT, DT_BFLOAT16})
.TypeConstraint("T", {DT_DOUBLE, DT_FLOAT, DT_BFLOAT16})
.TypeConstraint("Tseed", DT_INT32),
StatelessCategoricalOp);