From bace8b06d0a6e80ef1f0d6cf20c9b3ae78fc9362 Mon Sep 17 00:00:00 2001
From: Frederic Bastien <fbastien@nvidia.com>
Date: Tue, 8 Oct 2019 12:28:47 -0700
Subject: [PATCH] Add utility fct FillWithZY

---
 tensorflow/compiler/xla/array4d.h | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tensorflow/compiler/xla/array4d.h b/tensorflow/compiler/xla/array4d.h
index e23d317baf9..303640b654f 100644
--- a/tensorflow/compiler/xla/array4d.h
+++ b/tensorflow/compiler/xla/array4d.h
@@ -120,6 +120,21 @@ class Array4D : public Array<T> {
     }
   }
 
+  // Fills all of the {p,x} with the array provided, which specifies {z,y}.
+  void FillWithZY(const Array2D<T>& value) {
+    CHECK_EQ(value.height(), depth());
+    CHECK_EQ(value.width(), height());
+    for (int64 plane = 0; plane < planes(); ++plane) {
+      for (int64 depth = 0; depth < this->depth(); ++depth) {
+        for (int64 height = 0; height < this->height(); ++height) {
+          for (int64 width = 0; width < this->width(); ++width) {
+            (*this)(plane, depth, height, width) = value(depth, height);
+          }
+        }
+      }
+    }
+  }
+
   // Fills all of the {x,y} with the array provided, which specifies {p,z}.
   void FillWithPZ(const Array2D<T>& value) {
     CHECK_EQ(value.height(), planes());