diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2017-08-16 16:19:33 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2017-08-16 16:23:24 -0700 |
| commit | 018530bbb673ef630d9235db436a59f66b876d91 (patch) | |
| tree | 5d208d33461c08270e3bd935cca1205aec89a5cd /tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc | |
| parent | 3292c1a7369abbed73249b43432abbb0f86b4bec (diff) | |
Use XLA While loop in Gather/Scatter tf2xla bridge ops, to prevent excessively large HLO output for models with large numbers of embedding lookups.
PiperOrigin-RevId: 165514289
Diffstat (limited to 'tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc')
| -rw-r--r-- | tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc index 8b9ef4de76..bdd7e73302 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc @@ -312,7 +312,7 @@ class TensorArrayGatherOp : public XlaOpKernel { xla::ComputationDataHandle ta = resource->value; xla::ComputationDataHandle gather = XlaComputeGatherDynamicSlice( - ta, ta_shape, indices, indices_shape, dtype_, b); + ctx, ta, ta_shape, indices, indices_shape, dtype_, b); ctx->SetOutput(0, gather); } |