diff options
| author |  Michael Kuperstein <mkuper@google.com> | 2018-07-10 13:04:30 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-10 13:08:18 -0700 |
| commit | 1882427291f212cd04d09f2b35af4a78f6d771b5 (patch) | |
| tree | c6172cec133d7f7fed37757db67822e074d658a9 /tensorflow/compiler/xla/service/hlo_parser_test.cc | |
| parent | ccbbd4484a29fb3ee7d0d67abcebafdb48c9059b (diff) | |
[XLA] Generalize sort semantics to Rk.
PiperOrigin-RevId: 203997296
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_parser_test.cc')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_parser_test.cc | 29 |
1 files changed, 27 insertions, 2 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_parser_test.cc b/tensorflow/compiler/xla/service/hlo_parser_test.cc index 88f3309baa..f06c705c42 100644 --- a/tensorflow/compiler/xla/service/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/service/hlo_parser_test.cc @@ -840,7 +840,7 @@ R"(HloModule sort ENTRY Sort { x = f32[1024]{0} parameter(0) - ROOT sorted = f32[1024]{0} sort(x) + ROOT sorted = f32[1024]{0} sort(x), dimensions={0} } )" @@ -853,7 +853,32 @@ R"(HloModule sort ENTRY Sort { keys = f32[1024]{0} parameter(0) values = s32[1024]{0} parameter(1) - ROOT sorted = (f32[1024]{0}, s32[1024]{0}) sort(keys, values) + ROOT sorted = (f32[1024]{0}, s32[1024]{0}) sort(keys, values), dimensions={0} +} + +)" +}, +// R2 Sort (Key) +{ +"SortKeyR2", +R"(HloModule sort + +ENTRY Sort { + x = f32[1024,16]{0,1} parameter(0) + ROOT sorted = f32[1024,16]{0,1} sort(x), dimensions={0} +} + +)" +}, +// R2 Sort (Key, Value) +{ +"SortKeyValueR2", +R"(HloModule sort + +ENTRY Sort { + keys = f32[1024,16]{0,1} parameter(0) + values = s32[1024,16]{0,1} parameter(1) + ROOT sorted = (f32[1024,16]{0,1}, s32[1024,16]{0,1}) sort(keys, values), dimensions={0} } )" |