From f947e4d46721306a0b68a58fc8ea74a102f0d041 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 8 Jan 2020 06:57:07 -0800
Subject: [PATCH] MaxUnpooling3DAttributes and functions for shape calculation.

PiperOrigin-RevId: 288687596
Change-Id: Icc3ebc21c564f88c89dbc9b203cc738c0def6e08
---
 .../lite/delegates/gpu/common/operations.cc   | 24 +++++++++++++++++++
 .../lite/delegates/gpu/common/operations.h    | 17 +++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/tensorflow/lite/delegates/gpu/common/operations.cc b/tensorflow/lite/delegates/gpu/common/operations.cc
index 7f884bcb0505a4..9002da4cf26c59 100644
--- a/tensorflow/lite/delegates/gpu/common/operations.cc
+++ b/tensorflow/lite/delegates/gpu/common/operations.cc
@@ -325,6 +325,13 @@ int32_t CalculateSamePadding(const BHWC& input,
                               /*dilation=*/1, attr.strides.get<AxisT>());
 }
 
+template <Axis AxisT>
+int32_t CalculateSamePadding(const BHWDC& input,
+                             const MaxUnpooling3DAttributes& attr) {
+  return CalculateSamePadding(input.get<AxisT>(), attr.kernel.get<AxisT>(),
+                              /*dilation=*/1, attr.strides.get<AxisT>());
+}
+
 Padding2D MakeSamePadding(const BHWC& input,
                           const ConvolutionTransposedAttributes& attr) {
   int32_t padding_height = CalculateSamePadding<Axis::HEIGHT>(input, attr);
@@ -375,6 +382,18 @@ BHWC CalculateOutputShape(const BHWC& input,
               input.c);
 }
 
+BHWDC CalculateOutputShape(const BHWDC& input,
+                           const MaxUnpooling3DAttributes& attr) {
+  return BHWDC(input.b,
+               input.h * attr.strides.h - attr.padding.prepended.h -
+                   attr.padding.appended.h,
+               input.w * attr.strides.w - attr.padding.prepended.w -
+                   attr.padding.appended.w,
+               input.d * attr.strides.d - attr.padding.prepended.d -
+                   attr.padding.appended.d,
+               input.c);
+}
+
 BHWC CalculateOutputShape(const BHWC& input, const Pooling2DAttributes& attr) {
   return BHWC(input.b, CalculateOutput<Axis::HEIGHT>(input, attr),
               CalculateOutput<Axis::WIDTH>(input, attr), input.c);
@@ -527,6 +546,11 @@ Padding2D CalculateSamePadding(const BHWC& input,
   return MakeSamePadding(input, attr);
 }
 
+Padding3D CalculateSamePadding(const BHWDC& input,
+                               const MaxUnpooling3DAttributes& attr) {
+  return MakeSamePadding(input, attr);
+}
+
 float CalculateResizeScale(int32_t input_size, int32_t output_size,
                            const Upsample2DAttributes& attr) {
   return attr.align_corners && input_size > 1 && output_size > 1
diff --git a/tensorflow/lite/delegates/gpu/common/operations.h b/tensorflow/lite/delegates/gpu/common/operations.h
index 5187e4192bb261..3ee375175a4fa2 100644
--- a/tensorflow/lite/delegates/gpu/common/operations.h
+++ b/tensorflow/lite/delegates/gpu/common/operations.h
@@ -159,6 +159,13 @@ struct MaxUnpooling2DAttributes {
   Padding2D padding;
 };
 
+struct MaxUnpooling3DAttributes {
+  // Strides for every axis.
+  HWD strides = HWD(0, 0, 0);
+  HWD kernel = HWD(0, 0, 0);
+  Padding3D padding;
+};
+
 struct ConcatAttributes {
   // Defines axis by which to concat on.
   Axis axis = Axis::UNKNOWN;
@@ -169,6 +176,11 @@ struct ConcatAttributes {
 BHWC CalculateOutputShape(const BHWC& input,
                           const MaxUnpooling2DAttributes& attr);
 
+// @return shape of a tensor after MaxUnpooling3D operation is applied to
+//         the given input.
+BHWDC CalculateOutputShape(const BHWDC& input,
+                           const MaxUnpooling3DAttributes& attr);
+
 // @return shape of a tensor after Pooling2D operation is applied to the given
 //         input.
 BHWC CalculateOutputShape(const BHWC& input, const Pooling2DAttributes& attr);
@@ -197,6 +209,11 @@ Padding3D CalculateSamePadding(const BHWDC& input,
 Padding2D CalculateSamePadding(const BHWC& input,
                                const MaxUnpooling2DAttributes& attr);
 
+// @return padding for max unpooling operation to make sure output keep the same
+// shape as the given input.
+Padding3D CalculateSamePadding(const BHWDC& input,
+                               const MaxUnpooling3DAttributes& attr);
+
 struct Convolution2DAttributes {
   HW strides = HW(1, 1);    // Along each axis.
   HW dilations = HW(1, 1);  // Along each axis.