diff options
| author |  David Majnemer <majnemer@google.com> | 2018-08-28 13:56:25 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-08-28 14:01:14 -0700 |
| commit | 96de2f020130fc0afbcc5087b520aa46ced5b5af (patch) | |
| tree | 759a2b4d58448a1cd816a3b84979dc6d40d36467 /tensorflow/compiler/xla/service/hlo_instruction.h | |
| parent | 9c1f14322484e44a93b77619ffd2e24b9b7a9b1d (diff) | |
[XLA] Implement kIota for CPU & GPU, extend it w/ broadcast semantics
This extends the Iota HLO to have a broadcast field. This allows for higher
rank kIota operations.
PiperOrigin-RevId: 210600435
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_instruction.h')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_instruction.h | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index b393635e9d..4a424cebc0 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -350,7 +350,8 @@ class HloInstruction { std::unique_ptr<Literal> literal); // Creates an Iota instruction. - static std::unique_ptr<HloInstruction> CreateIota(const Shape& shape); + static std::unique_ptr<HloInstruction> CreateIota(const Shape& shape, + int64 iota_dimension); // Creates a get tuple element instruction. static std::unique_ptr<HloInstruction> CreateGetTupleElement( |