diff options
| author |  Yifei Feng <yifeif@google.com> | 2018-04-23 21:19:14 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-04-23 21:21:38 -0700 |
| commit | 22f3a97b8b089202f60bb0c7697feb0c8e0713cc (patch) | |
| tree | d16f95826e4be15bbb3b0f22bed0ca25d3eb5897 /tensorflow/core/ops/random_ops.cc | |
| parent | 24b7c9a800ab5086d45a7d83ebcd6218424dc9e3 (diff) | |
Merge changes from github.
PiperOrigin-RevId: 194031845
Diffstat (limited to 'tensorflow/core/ops/random_ops.cc')
| -rw-r--r-- | tensorflow/core/ops/random_ops.cc | 7 |
1 files changed, 6 insertions, 1 deletions
diff --git a/tensorflow/core/ops/random_ops.cc b/tensorflow/core/ops/random_ops.cc index f6c668f5c9..416ce9c0d8 100644 --- a/tensorflow/core/ops/random_ops.cc +++ b/tensorflow/core/ops/random_ops.cc @@ -43,7 +43,12 @@ REGISTER_OP("RandomUniformInt") .Attr("seed2: int = 0") .Attr("Tout: {int32, int64}") .Attr("T: {int32, int64}") - .SetShapeFn(shape_inference::RandomShape); + .SetShapeFn([](InferenceContext* c) { + ShapeHandle unused; + TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused)); + TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused)); + return shape_inference::RandomShape(c); + }); REGISTER_OP("RandomStandardNormal") .Input("shape: T") |