diff options
| author |  Michael Kuperstein <mkuper@google.com> | 2018-06-29 10:57:24 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-06-29 10:59:52 -0700 |
| commit | ec0a702ff1f22b73cec2d8f14c7a84c5a02856fd (patch) | |
| tree | bbe4a56f06e157fb69628fd40e4ab66f0ece27ea /tensorflow/compiler/xla/service/hlo_instruction.h | |
| parent | aa060bebb0fb064460ae4c3e92a0272be4ea04de (diff) | |
[XLA] Add key-value version of Sort HLO.
This is only currently implemented in the evaluator backend, and even that implementation is partial - the key and value type must match.
PiperOrigin-RevId: 202673122
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_instruction.h')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_instruction.h | 5 |
1 files changed, 5 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index 59a383218c..0459072127 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -611,6 +611,11 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, tensorflow::gtl::ArraySlice<int64> dimensions); + // Creates a sort op, with a keys operand, and an optional values operand. + static std::unique_ptr<HloInstruction> CreateSort( + const Shape& shape, HloInstruction* keys, + HloInstruction* values = nullptr); + // Creates a while instruction, given a condition computation, a body // computation, and the initial value for the input of the computations. For // example, shape: S32, condition: i -> i < 1000, body: i -> i * 2, init: 1 |