diff options
| author |  Peter Hawkins <phawkins@google.com> | 2018-02-05 14:12:02 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-02-05 14:16:03 -0800 |
| commit | 473bc3580510e2da299662da200791cf4a9fb086 (patch) | |
| tree | c0993c14233818f1cddee862781d077c0c5c9961 /tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc | |
| parent | 1bbfc0c9cebd1808fa46024f738e56d10d57d97e (diff) | |
[TF:XLA] Implement GatherNd.
PiperOrigin-RevId: 184584104
Diffstat (limited to 'tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc')
| -rw-r--r-- | tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc | 7 |
1 files changed, 5 insertions, 2 deletions
diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc index 7cf9b796b9..000b50af6b 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc @@ -337,8 +337,11 @@ class TensorArrayGatherOp : public XlaOpKernel { } } - xla::ComputationDataHandle gather = XlaComputeGatherDynamicSlice( - ctx, ta, ta_shape, indices, indices_shape, 0, dtype_, index_type, b); + xla::ComputationDataHandle gather; + OP_REQUIRES_OK( + ctx, + XlaGather(ta, ta_shape, indices, indices_shape, /*axis=*/0, + /*indices_are_nd=*/false, dtype_, index_type, b, &gather)); ctx->SetOutput(0, gather); } |