Support SpaceToBatchND in TFLite.

The internal implementation only supports 4D tensors for now.
The dimension has to be 1 batch + 2 spatial + 1 other.
The most common format within this restriction is NHWC.

PiperOrigin-RevId: 180592870

1 files changed, 11 insertions, 0 deletions

diff --git a/tensorflow/contrib/lite/builtin_op_data.h b/tensorflow/contrib/lite/builtin_op_data.h
index 5c6f3016b1..0c6e7f93a6 100644
--- a/tensorflow/contrib/lite/builtin_op_data.h
+++ b/tensorflow/contrib/lite/builtin_op_data.h
@@ -111,6 +111,17 @@ typedef struct {
   // TODO(ahentz): We can't have dynamic data in this struct, at least not yet.
   // For now we will fix the maximum possible number of dimensions.
   int block_shape[2];
+  int before_paddings[2];
+  int after_paddings[2];
+} TfLiteSpaceToBatchNDParams;
+
+typedef struct {
+  // Number of spatial dimensions.
+  // For now only NHWC is supported, and the value should always be 2.
+  int num_spatial_dimensions;
+  // TODO(ahentz): We can't have dynamic data in this struct, at least not yet.
+  // For now we will fix the maximum possible number of dimensions.
+  int block_shape[2];
   int before_crops[2];
   int after_crops[2];
 } TfLiteBatchToSpaceNDParams;