Force multislice/update to be 2D

Summary: Ref T45278

Test Plan: CI

Reviewers: jakeh, georgew, jackh, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved

Reviewed By: georgew, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved

Maniphest Tasks: T45278

Differential Revision: https://phabricator.sourcevertex.net/D52785

Author: georgepaw

tensorflow/compiler/plugin/poplar/kernels/popops/multi_slice.cc

@@ -30,7 +30,17 @@ class PopopsMultiSliceOp : public XlaOpKernel, IpuOpKernel {
 
   void Compile(XlaOpKernelContext* ctx) override {
     const TensorShape input_shape = ctx->InputShape(0);
-    TensorShape output_shape = ctx->InputShape(1);
+    const TensorShape indices_shape = ctx->InputShape(1);
+
+    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(input_shape),
+                errors::InvalidArgument("input shape must be 2D, but got: ",
+                                        input_shape.DebugString()));
+
+    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(indices_shape),
+                errors::InvalidArgument("indices shape must be 1D, but got: ",
+                                        indices_shape.DebugString()));
+
+    TensorShape output_shape = indices_shape;
     output_shape.AddDim(input_shape.dim_size(1));
 
     xla::PrimitiveType input_type;
@@ -63,10 +73,18 @@ class PopopsMultiUpdateOp : public XlaOpKernel, IpuOpKernel {
     const TensorShape input_shape = ctx->InputShape(0);
     const TensorShape indices_shape = ctx->InputShape(1);
     const TensorShape updates_shape = ctx->InputShape(2);
-    xla::PrimitiveType input_type;
-    OP_REQUIRES_OK(ctx,
-                   DataTypeToPrimitiveType(ctx->input_type(0), &input_type));
 
+    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(input_shape),
+                errors::InvalidArgument("input shape must be 2D, but got: ",
+                                        input_shape.DebugString()));
+
+    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(indices_shape),
+                errors::InvalidArgument("indices shape must be 1D, but got: ",
+                                        indices_shape.DebugString()));
+
+    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(updates_shape),
+                errors::InvalidArgument("updates shape must be 2D, but got: ",
+                                        updates_shape.DebugString()));
     if (is_update_add_) {
       const TensorShape scale_shape = ctx->InputShape(3);
       OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(scale_shape),
@@ -75,6 +93,10 @@ class PopopsMultiUpdateOp : public XlaOpKernel, IpuOpKernel {
     }
 
     xla::XlaBuilder& b = *ctx->builder();
+
+    xla::PrimitiveType input_type;
+    OP_REQUIRES_OK(ctx,
+                   DataTypeToPrimitiveType(ctx->input_type(0), &input_type));
     xla::Shape xla_output_shape =
         TensorShapeToXLAShape(input_type, input_shape);
     const auto num_inputs = ctx->num_inputs();

tensorflow/python/ipu/tests/embedding_lookup_test.py

@@ -204,8 +204,8 @@ def network(x1, lr):
   def testGradient(self):
     with self.session() as sess:
       with ops.device('cpu'):
-        x1 = array_ops.placeholder(np.int32, shape=[3, 4, 2])
-        grads = array_ops.placeholder(np.float32, shape=[3, 4, 2, 16])
+        x1 = array_ops.placeholder(np.int32, shape=[24])
+        grads = array_ops.placeholder(np.float32, shape=[24, 16])
         lr = array_ops.placeholder(np.float32, shape=[])
 
       def network(x1, grads, lr):
@@ -231,9 +231,10 @@ def network(x1, grads, lr):
       sess.run(variables.global_variables_initializer())
       out, indices, gradient = sess.run(
           r, {
-              x1: [[[10, 11], [12, 13], [14, 15], [16, 17]],
-                   [[20, 21], [22, 23], [24, 25], [26, 27]],
-                   [[30, 31], [32, 33], [34, 35], [36, 37]]],
+              x1: [
+                  10, 11, 12, 13, 14, 15, 16, 17, 20, 21, 22, 23, 24, 25, 26,
+                  27, 30, 31, 32, 33, 34, 35, 36, 37
+              ],
               grads:
               np.random.rand(*grads.shape),
               lr:

tensorflow/python/ipu/tests/replication_all_to_all_test.py

@@ -77,7 +77,7 @@ def testAllGatherShapeInference(self):
   @test_util.deprecated_graph_mode_only
   def testSerializedMultiUpdateAdd(self):
     with ops.device('cpu'):
-      idx = array_ops.placeholder(np.int32, shape=[16, 1])
+      idx = array_ops.placeholder(np.int32, shape=[16])
       updates = array_ops.placeholder(np.float32, shape=[16, 128])
       scale = array_ops.placeholder(np.float32, shape=[])
 
@@ -108,12 +108,10 @@ def my_graph(idx, updates, scale):
       sess.run(variables.global_variables_initializer())
       sess.run(
           out, {
-              idx: [[1], [2], [3], [4], [1], [2], [3], [4], [10], [20], [30],
-                    [40], [100], [200], [300], [400]],
-              updates:
-              np.ones(updates.shape),
-              scale:
-              2,
+              idx:
+              [1, 2, 3, 4, 1, 2, 3, 4, 10, 20, 30, 40, 100, 200, 300, 400],
+              updates: np.ones(updates.shape),
+              scale: 2,
           })
       result = sess.run(outfeed)