diff options
author | Mark Heffernan <meheff@google.com> | 2018-10-08 14:26:43 -0700 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-10-08 14:34:02 -0700 |
commit | 396a8a4105edd409d0821c4d5d0b920b315ffb72 (patch) | |
tree | 428350d427ffb29470e284077a2734b273b7cc4d /tensorflow/compiler/xla/layout_util.cc | |
parent | bc5635dc3ac78007caee88fabd81d23ad945b637 (diff) |
Add custom call with layout constraints.
Add a variant of CustomCall which specifies arbitrary layout constraints on the operands and result. The existing non-layout-constrained CustomCall is changed to have no layout preference and can now be assigned arbitrary layouts by layout assignment.
PiperOrigin-RevId: 216249615
Diffstat (limited to 'tensorflow/compiler/xla/layout_util.cc')
-rw-r--r-- | tensorflow/compiler/xla/layout_util.cc | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/layout_util.cc b/tensorflow/compiler/xla/layout_util.cc index d310335618..3c8db9aa45 100644 --- a/tensorflow/compiler/xla/layout_util.cc +++ b/tensorflow/compiler/xla/layout_util.cc @@ -65,6 +65,12 @@ void SetDefaultLayoutToContainer( return layout; } +/* static */ Layout LayoutUtil::MakeDescendingLayout(int64 rank) { + std::vector<int64> layout(rank); + std::iota(layout.rbegin(), layout.rend(), static_cast<int64>(0)); + return MakeLayout(layout); +} + /* static */ Layout LayoutUtil::MakeLayoutFromMajorToMinor( absl::Span<const int64> major_to_minor) { Layout layout; |