Incomplete validation of shapes in multiple TF ops

@@ -183,20 +183,29 @@ struct LaunchGrouped {
     auto on_shuffled = [&]() { shuffles_completed.DecrementCount(); };
 
     // Shuffle input into temporary tensor.
-    Tensor input_shuffled(input.dtype(), TensorShape(post_shuffle(input)));
+    Tensor input_shuffled;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_temp(input.dtype(), TensorShape(post_shuffle(input)),
+                                &input_shuffled));
     input_shuffled.tensor<T, 5>().device(device, on_shuffled) =
         input.shaped<T, 5>(pre_shuffle(input)).shuffle(shuffle);
 
     // Shuffle filter into temporary tensor.
-    Tensor filter_shuffled(filter.dtype(), TensorShape(post_shuffle(filter)));
+    Tensor filter_shuffled;
+    OP_REQUIRES_OK(ctx, ctx->allocate_temp(filter.dtype(),
+                                           TensorShape(post_shuffle(filter)),
+                                           &filter_shuffled));
     filter_shuffled.tensor<T, 5>().device(device, on_shuffled) =
         filter.shaped<T, 5>(pre_shuffle(filter)).shuffle(shuffle);
 
     // Wait for the completion of input/filter shuffles.
     shuffles_completed.Wait();
 
     // Write group convolution results into temporary output tensor.
-    Tensor output_shuffled(output->dtype(), TensorShape(post_shuffle(*output)));
+    Tensor output_shuffled;
+    OP_REQUIRES_OK(ctx, ctx->allocate_temp(output->dtype(),
+                                           TensorShape(post_shuffle(*output)),
+                                           &output_shuffled));
 
     for (int64_t i = 0; i < num_groups; ++i) {
       // TODO(ezhulenev): Run this loop using `parallelFor` (regular parallelFor

@@ -189,6 +189,9 @@ class MaxPoolingOp : public OpKernel {
   void SpatialMaxPool(OpKernelContext* context, Tensor* output,
                       const Tensor& tensor_in, const PoolParameters& params,
                       const Padding& padding) {
+    if (output->NumElements() == 0) {
+      return;
+    }
     // On GPU, use Eigen's Spatial Max Pooling.  On CPU, use an
     // EigenMatrix version that is currently faster than Eigen's
     // Spatial MaxPooling implementation.
@@ -443,6 +446,9 @@ class MaxPoolingV2Op : public OpKernel {
   void SpatialMaxPool(OpKernelContext* context, Tensor* output,
                       const Tensor& tensor_in, const PoolParameters& params,
                       const Padding& padding) {
+    if (output->NumElements() == 0) {
+      return;
+    }
     // On GPU, use Eigen's Spatial Max Pooling.  On CPU, use an
     // EigenMatrix version that is currently faster than Eigen's
     // Spatial MaxPooling implementation.
@@ -561,6 +567,9 @@ template <typename Device, typename T>
 void SpatialAvgPool(OpKernelContext* context, Tensor* output,
                     const Tensor& input, const PoolParameters& params,
                     const Padding& padding) {
+  if (output->NumElements() == 0) {
+    return;
+  }
   typedef Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
       ConstEigenMatrixMap;
   typedef Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>

@@ -325,6 +325,14 @@ class MaxPoolingGradOp : public OpKernel {
     if (!context->status().ok()) {
       return;
     }
+    OP_REQUIRES(context, tensor_out.shape() == params.forward_output_shape(),
+                errors::InvalidArgument("Expected orig_output shape to be ",
+                                        params.forward_output_shape(),
+                                        ", but got ", tensor_out.shape()));
+    OP_REQUIRES(context, out_backprop.shape() == params.forward_output_shape(),
+                errors::InvalidArgument("Expected grad shape to be ",
+                                        params.forward_output_shape(),
+                                        ", but got ", out_backprop.shape()));
 
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
@@ -538,6 +546,18 @@ class MaxPoolingGradGradOp : public OpKernel {
                           /*explicit_paddings=*/{},
                           FORMAT_NHWC,
                           tensor_in.shape()};
+    if (!context->status().ok()) {
+      return;
+    }
+    OP_REQUIRES(context, tensor_out.shape() == params.forward_output_shape(),
+                errors::InvalidArgument("Expected orig_output shape to be ",
+                                        params.forward_output_shape(),
+                                        ", but got ", tensor_out.shape()));
+    OP_REQUIRES(
+        context, out_grad_backprop.shape() == tensor_in.shape(),
+        errors::InvalidArgument("Expected grad shape to be ", tensor_in.shape(),
+                                ", but got ", out_grad_backprop.shape()));
+
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
                                 {2}, 0, tensor_out.shape(), &output));
@@ -742,6 +762,17 @@ class MaxPoolingGradGradOp<Eigen::GpuDevice, T> : public OpKernel {
                           /*explicit_paddings=*/{},
                           data_format_,
                           tensor_in.shape()};
+    if (!context->status().ok()) {
+      return;
+    }
+    OP_REQUIRES(context, tensor_out.shape() == params.forward_output_shape(),
+                errors::InvalidArgument("Expected orig_output shape to be ",
+                                        params.forward_output_shape(),
+                                        ", but got ", tensor_out.shape()));
+    OP_REQUIRES(
+        context, out_grad_backprop.shape() == tensor_in.shape(),
+        errors::InvalidArgument("Expected grad shape to be ", tensor_in.shape(),
+                                ", but got ", out_grad_backprop.shape()));
 
     functor::MaxPoolGradBackwardNoMask<T>()(
         data_format_, tensor_in.flat<T>().data(), tensor_out.flat<T>().data(),
@@ -1096,6 +1127,14 @@ class MaxPoolingGradWithArgmaxOp : public OpKernel {
     if (!context->status().ok()) {
       return;
     }
+    OP_REQUIRES(context, grad_in.shape() == params.forward_output_shape(),
+                errors::InvalidArgument("Expected grad shape to be ",
+                                        params.forward_output_shape(),
+                                        ", but got ", grad_in.shape()));
+    OP_REQUIRES(context, argmax.shape() == params.forward_output_shape(),
+                errors::InvalidArgument("Expected argmax shape to be ",
+                                        params.forward_output_shape(),
+                                        ", but got ", argmax.shape()));
 
     TensorShape out_shape({params.tensor_in_batch, params.tensor_in_rows,
                            params.tensor_in_cols, params.depth});
@@ -1156,6 +1195,14 @@ class MaxPoolingGradGradWithArgmaxOp : public OpKernel {
     if (!context->status().ok()) {
       return;
     }
+    OP_REQUIRES(
+        context, grad_in.shape() == tensor_in.shape(),
+        errors::InvalidArgument("Expected grad shape to be ", tensor_in.shape(),
+                                ", but got ", grad_in.shape()));
+    OP_REQUIRES(context, argmax.shape() == params.forward_output_shape(),
+                errors::InvalidArgument("Expected argmax shape to be ",
+                                        params.forward_output_shape(),
+                                        ", but got ", argmax.shape()));
 
     TensorShape out_shape({params.tensor_in_batch, params.out_height,
                            params.out_width, params.depth});

@@ -366,6 +366,19 @@ class MaxPooling3dGradOp : public OpKernel {
 
     OP_REQUIRES_OK(context, Get3dOutputSize(input_size, window, stride,
                                             padding_, &out, &padding));
+
+    const int64_t depth = GetTensorDim(tensor_in, data_format_, 'C');
+    const int64_t in_batch = GetTensorDim(tensor_in, data_format_, 'N');
+    TensorShape out_shape = ShapeFromFormat(data_format_, in_batch,
+                                            {{out[2], out[1], out[0]}}, depth);
+    OP_REQUIRES(
+        context, tensor_out.shape() == out_shape,
+        errors::InvalidArgument("Expected orig_output shape to be ", out_shape,
+                                ", but got ", tensor_out.shape()));
+    OP_REQUIRES(context, out_backprop.shape() == out_shape,
+                errors::InvalidArgument("Expected grad shape to be ", out_shape,
+                                        ", but got ", out_backprop.shape()));
+
     LaunchMaxPooling3dGradOp<Device, T>::launch(
         context, tensor_in, tensor_out, out_backprop, window, stride, out,
         padding, data_format_, input_backprop);
@@ -712,6 +725,14 @@ class MaxPooling3dGradGradOp : public OpKernel {
     Pool3dParameters params{context,  ksize_,       stride_,
                             padding_, data_format_, tensor_in.shape()};
     if (!context->status().ok()) return;  // params is invalid
+    OP_REQUIRES(context, tensor_out.shape() == params.forward_output_shape(),
+                errors::InvalidArgument("Expected orig_output shape to be ",
+                                        params.forward_output_shape(),
+                                        ", but got ", tensor_out.shape()));
+    OP_REQUIRES(
+        context, out_grad_backprop.shape() == tensor_in.shape(),
+        errors::InvalidArgument("Expected grad shape to be ", tensor_in.shape(),
+                                ", but got ", out_grad_backprop.shape()));
 
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(

@@ -465,6 +465,16 @@ void DnnPoolingGradOp<T>::Compute(
   if (!context->status().ok()) {
     return;
   }
+  if (tensor_out) {
+    OP_REQUIRES(context, tensor_out->shape() == params.forward_output_shape(),
+                errors::InvalidArgument("Expected orig_output shape to be ",
+                                        params.forward_output_shape(),
+                                        ", but got ", tensor_out->shape()));
+  }
+  OP_REQUIRES(context, out_backprop.shape() == params.forward_output_shape(),
+              errors::InvalidArgument("Expected grad shape to be ",
+                                      params.forward_output_shape(),
+                                      ", but got ", out_backprop.shape()));
 
   TensorFormat transformed_input_data_format = data_format;
 

@@ -83,11 +83,6 @@ struct PoolParameters {
   TensorFormat data_format;
 };
 
-// Checks if the sizes of the paddings are less than the size of window.
-// This is required for MaxPool because it pads with -inf, so the pooling
-// window cannot fully cover the padded area.
-Status CheckPaddingSize(PoolParameters& params);
-
 // An implementation of MaxPooling (forward).
 // TODO (yongtang): Remove MaxPoolingOp and use MaxPoolingV2Op,
 //     QuantizedMaxPoolingOp depends on MaxPoolingOp so keep intact for now

@@ -16,9 +16,13 @@
 
 import numpy as np
 
+from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import errors
+from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
 from tensorflow.python.ops import gradient_checker
 from tensorflow.python.ops import gradients_impl
 from tensorflow.python.ops import nn_ops
@@ -515,6 +519,44 @@ def testMaxPool3DZeroPoolSize(self):
           pool_3d = f(input_tensor, ksize=[2, 2, 0], strides=1, padding="VALID")
           self.evaluate(pool_3d)
 
+  def testMaxPoolGradEagerShapeErrors(self):
+    with context.eager_mode():
+      orig_in = array_ops.ones((1, 1, 1, 1, 1))
+
+      # Test invalid orig_out shape
+      orig_out = array_ops.ones((1, 1, 1, 1, 2))
+      grad = array_ops.ones((1, 1, 1, 1, 1))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected orig_output shape to be \[1,1,1,1,1\], but got "
+          r"\[1,1,1,1,2\]"):
+        gen_nn_ops.max_pool3d_grad(
+            orig_in, orig_out, grad, ksize=[1, 1, 1, 1, 1],
+            strides=[1, 1, 1, 1, 1], padding="VALID")
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected orig_output shape to be \[1,1,1,1,1\], but got "
+          r"\[1,1,1,1,2\]"):
+        gen_nn_ops.max_pool3d_grad_grad(
+            orig_in, orig_out, grad, ksize=[1, 1, 1, 1, 1],
+            strides=[1, 1, 1, 1, 1], padding="VALID")
+
+      # Test invalid grad shape
+      orig_out = array_ops.ones((1, 1, 1, 1, 1))
+      grad = array_ops.ones((1, 1, 1, 1, 2))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected grad shape to be \[1,1,1,1,1\], but got \[1,1,1,1,2\]"):
+        gen_nn_ops.max_pool3d_grad(
+            orig_in, orig_out, grad, ksize=[1, 1, 1, 1, 1],
+            strides=[1, 1, 1, 1, 1], padding="VALID")
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected grad shape to be \[1,1,1,1,1\], but got \[1,1,1,1,2\]"):
+        gen_nn_ops.max_pool3d_grad_grad(
+            orig_in, orig_out, grad, ksize=[1, 1, 1, 1, 1],
+            strides=[1, 1, 1, 1, 1], padding="VALID")
+
 
 if __name__ == "__main__":
   test.main()

@@ -618,6 +618,7 @@ def testMaxPoolExplicitPaddingAdvanced(self, **kwargs):
 
   @parameterized.parameters(
       GetTestConfigsDicts(nn_ops.max_pool, nn_ops.max_pool_v2))
+  @test_util.xla_allow_fallback("XLA doesn't support explicit padding")
   @test_util.run_deprecated_v1
   def testMaxPoolNegativeInputExpPaddingAdv(self, **kwargs):
     expected_output = [-1, -1, -3, -5, -7, -7, -9, -11, -19, -19, -21, -23, -31,
@@ -2390,6 +2391,82 @@ def testExplicitPaddingBatch(self):
             explicit_paddings=[1, 1, 1, 1, 1, 1, 0, 0],
             data_format="NHWC"))
 
+  def testMaxPoolGradEagerShapeErrors(self):
+    with context.eager_mode():
+      orig_in = array_ops.ones((1, 1, 1, 1))
+
+      # Test invalid orig_out shape
+      orig_out = array_ops.ones((1, 1, 1, 2))
+      grad = array_ops.ones((1, 1, 1, 1))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected orig_output shape to be \[1,1,1,1\], but got \[1,1,1,2\]"):
+        gen_nn_ops.max_pool_grad(
+            orig_in, orig_out, grad, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
+            padding="VALID")
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected orig_output shape to be \[1,1,1,1\], but got \[1,1,1,2\]"):
+        gen_nn_ops.max_pool_grad_grad(
+            orig_in, orig_out, grad, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
+            padding="VALID")
+
+      # Test invalid grad shape
+      orig_out = array_ops.ones((1, 1, 1, 1))
+      grad = array_ops.ones((1, 1, 1, 2))
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected grad shape to be \[1,1,1,1\], but got \[1,1,1,2\]"):
+        gen_nn_ops.max_pool_grad(
+            orig_in, orig_out, grad, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
+            padding="VALID")
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected grad shape to be \[1,1,1,1\], but got \[1,1,1,2\]"):
+        gen_nn_ops.max_pool_grad_grad(
+            orig_in, orig_out, grad, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
+            padding="VALID")
+
+  def testMaxPoolGradWithArgmaxEagerShapeErrors(self):
+    with context.eager_mode():
+      inp = array_ops.ones((1, 1, 1, 1))
+
+      # Test invalid grad shape
+      grad = array_ops.ones((1, 1, 1, 2))
+      argmax = array_ops.zeros((1, 1, 1, 1), dtype=dtypes.int64)
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected grad shape to be \[1,1,1,1\], but got \[1,1,1,2\]"):
+        gen_nn_ops.max_pool_grad_with_argmax(
+            inp, grad, argmax, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
+            padding="VALID")
+      # max_pool_grad_grad_with_argmax is only implemented for GPUs
+      if test.is_gpu_available():
+        with self.assertRaisesRegex(
+            errors_impl.InvalidArgumentError,
+            r"Expected grad shape to be \[1,1,1,1\], but got \[1,1,1,2\]"):
+          gen_nn_ops.max_pool_grad_grad_with_argmax(
+              inp, grad, argmax, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
+              padding="VALID")
+
+      # Test invalid argmax shape
+      grad = array_ops.ones((1, 1, 1, 1))
+      argmax = array_ops.ones((1, 1, 1, 2), dtype=dtypes.int64)
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          r"Expected argmax shape to be \[1,1,1,1\], but got \[1,1,1,2\]"):
+        gen_nn_ops.max_pool_grad_with_argmax(
+            inp, grad, argmax, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
+            padding="VALID")
+      # max_pool_grad_grad_with_argmax is only implemented for GPUs
+      if test.is_gpu_available():
+        with self.assertRaisesRegex(
+            errors_impl.InvalidArgumentError,
+            r"Expected argmax shape to be \[1,1,1,1\], but got \[1,1,1,2\]"):
+          gen_nn_ops.max_pool_grad_grad_with_argmax(
+              inp, grad, argmax, ksize=[1, 1, 1, 1], strides=[1, 1, 1, 1],
+              padding="VALID")
+
 
 def GetMaxPoolFwdTest(input_size, filter_size, strides, padding):
 

@@ -66,6 +66,12 @@ class TridiagonalMatMulOpGpu : public OpKernel {
     const Tensor& rhs = context->input(3);
 
     const int ndims = rhs.dims();
+    OP_REQUIRES(
+        context, ndims >= 2,
+        errors::InvalidArgument("Input must have rank >= 2, but got ", ndims));
+    OP_REQUIRES_OK(context, ValidateInputTensor(superdiag, "superdiag", rhs));
+    OP_REQUIRES_OK(context, ValidateInputTensor(maindiag, "maindiag", rhs));
+    OP_REQUIRES_OK(context, ValidateInputTensor(subdiag, "subdiag", rhs));
     int64 batch_size = 1;
     for (int i = 0; i < ndims - 2; i++) {
       batch_size *= rhs.dim_size(i);
@@ -85,6 +91,39 @@ class TridiagonalMatMulOpGpu : public OpKernel {
         maindiag.flat<Scalar>().data(), subdiag.flat<Scalar>().data(),
         rhs.flat<Scalar>().data(), output->flat<Scalar>().data()));
   }
+
+ private:
+  Status ValidateInputTensor(const Tensor& tensor,
+                             const std::string& tensor_name,
+                             const Tensor& rhs) {
+    const int ndims = rhs.dims();
+    if (tensor.dims() != ndims) {
+      return errors::InvalidArgument(tensor_name,
+                                     " must have same rank as rhs, but got ",
+                                     tensor.dims(), " and ", ndims);
+    }
+    for (int i = 0; i < ndims - 2; i++) {
+      if (tensor.dim_size(i) != rhs.dim_size(i)) {
+        return errors::InvalidArgument(
+            tensor_name,
+            " must have same outer dimensions as rhs, but for index ", i,
+            ", got ", tensor.dim_size(i), " and ", rhs.dim_size(i));
+      }
+    }
+    if (tensor.dim_size(ndims - 2) != 1) {
+      return errors::InvalidArgument(
+          tensor_name, "'s second-to-last dimension must be 1, but got ",
+          tensor.dim_size(ndims - 2));
+    }
+    if (tensor.dim_size(ndims - 1) != rhs.dim_size(ndims - 2)) {
+      return errors::InvalidArgument(tensor_name,
+                                     "'s last dimension size must be rhs's "
+                                     "second-to-last dimension size, but got ",
+                                     tensor.dim_size(ndims - 1), " and ",
+                                     rhs.dim_size(ndims - 2));
+    }
+    return Status::OK();
+  }
 };
 
 REGISTER_LINALG_OP_GPU("TridiagonalMatMul", (TridiagonalMatMulOpGpu<float>),

@@ -19,12 +19,15 @@
 import numpy as np
 
 from tensorflow.python.client import session
+from tensorflow.python.eager import context
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import gradient_checker_v2
+from tensorflow.python.ops import linalg_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import variables
 from tensorflow.python.ops.linalg import linalg_impl
@@ -175,6 +178,37 @@ def testGradientComplexWithBatches(self):
     rhs = self._randomComplexArray((b, m, n))
     self._gradientTest(diags, rhs, dtype=dtypes.complex128)
 
+  def _testErrorWithShapesEager(self, exception_regex, superdiag_shape,
+                                maindiag_shape, subdiag_shape, rhs_shape):
+    with context.eager_mode():
+      superdiag = array_ops.ones(superdiag_shape)
+      maindiag = array_ops.ones(maindiag_shape)
+      subdiag = array_ops.ones(subdiag_shape)
+      rhs = array_ops.ones(rhs_shape)
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  exception_regex):
+        linalg_ops.tridiagonal_mat_mul(superdiag, maindiag, subdiag, rhs)
+
+  def testInvalidShapesEagerGpu(self):
+    if not test.is_gpu_available():
+      self.skipTest('Test requires GPU')
+    self._testErrorWithShapesEager('Input must have rank >= 2, but got ',
+                                   [2], [2], [2], [2])
+    self._testErrorWithShapesEager(
+        'superdiag must have same rank as rhs, but got 3 and 2',
+        [2, 1, 2], [2, 1], [2, 1], [2, 2])
+    self._testErrorWithShapesEager(
+        'maindiag must have same outer dimensions as rhs, but for index 0, got '
+        '3 and 2',
+        [2, 1, 2], [3, 1, 2], [2, 1, 2], [2, 2, 2])
+    self._testErrorWithShapesEager(
+        "subdiag's second-to-last dimension must be 1, but got 3",
+        [2, 1, 2], [2, 1, 2], [2, 3, 2], [2, 2, 2])
+    self._testErrorWithShapesEager(
+        "subdiag's last dimension size must be rhs's second-to-last dimension "
+        "size, but got 3 and 2",
+        [2, 1, 2], [2, 1, 2], [2, 1, 3], [2, 2, 2])
+
   # Benchmark
   class TridiagonalMatMulBenchmark(test.Benchmark):
     sizes = [(100000, 1, 1), (1000000, 1, 1), (10000000, 1, 1), (100000, 10, 1),

@@ -46,7 +46,8 @@ class SoftmaxXentWithLogitsOp : public OpKernel {
     TensorShape shape_in = logits_in.shape();
 
     BCast bcast(BCast::FromShape(logits_in.shape()),
-                BCast::FromShape(labels_in.shape()));
+                BCast::FromShape(labels_in.shape()),
+                /*fewer_dims_optimization=*/false);
     if (!logits_in.IsSameSize(labels_in)) {
       OP_REQUIRES(context, bcast.IsValid(),
                   errors::InvalidArgument(
@@ -88,20 +89,12 @@ class SoftmaxXentWithLogitsOp : public OpKernel {
                                 {0}, 1, shape_in, &back_out));
     if (shape_in.dim_size(0) > 0) {
       functor::XentFunctor<Device, T> functor;
-      if (logits_in.IsSameSize(labels_in)) {
-        functor(context->eigen_device<Device>(), shape_in.AsEigenDSizes<2>(),
-                Eigen::array<Eigen::DenseIndex, 2>{1, 1},
-                Eigen::array<Eigen::DenseIndex, 2>{1, 1}, logits_in.matrix<T>(),
-                labels_in.matrix<T>(), scratch.matrix<T>(), loss_out->vec<T>(),
-                back_out->matrix<T>());
-      } else {
-        functor(context->eigen_device<Device>(), shape_in.AsEigenDSizes<2>(),
-                BCast::ToIndexArray<2>(bcast.x_bcast()),
-                BCast::ToIndexArray<2>(bcast.y_bcast()),
-                logits_in.template shaped<T, 2>(bcast.x_reshape()),
-                labels_in.template shaped<T, 2>(bcast.y_reshape()),
-                scratch.matrix<T>(), loss_out->vec<T>(), back_out->matrix<T>());
-      }
+      functor(context->eigen_device<Device>(), shape_in.AsEigenDSizes<2>(),
+              BCast::ToIndexArray<2>(bcast.x_bcast()),
+              BCast::ToIndexArray<2>(bcast.y_bcast()),
+              logits_in.template shaped<T, 2>(bcast.x_reshape()),
+              labels_in.template shaped<T, 2>(bcast.y_reshape()),
+              scratch.matrix<T>(), loss_out->vec<T>(), back_out->matrix<T>());
     }
   }
 };

@@ -151,6 +151,9 @@ def _testLabelsBroadcast(self, uniform_labels_gradient):
     labels = np.array([[0., 0., 0., 1.]]).astype(np.float16)
     logits = np.array([[1., 1., 1., 1.], [1., 2., 3., 4.]]).astype(np.float16)
     self._testXent2D(labels, logits, with_placeholders=True)
+    labels = np.array([[1.]]).astype(np.float16)
+    logits = np.array([[1.], [2.]]).astype(np.float16)
+    self._testXent2D(labels, logits, with_placeholders=True)
     labels = np.array([[0.], [2.], [0.25]]).astype(np.float16)
     logits = np.array([[1., 1., 1., 1.], [1., 2., 3., 4.],
                        [1., 2., 3., 4.]]).astype(np.float16)

@@ -63,6 +63,13 @@ def testFeaturesBroadcast(self):
     self.assertAllCloseAccordingToType(np_loss, tf_loss)
     self.assertAllCloseAccordingToType(np_gradient, tf_gradient)
 
+    tf_f = constant_op.constant(np.array([[1.]]).astype(np.float32))
+    tf_l = constant_op.constant(np.array([[1.], [1.]]).astype(np.float32))
+    tf_loss, tf_gradient = gen_nn_ops.softmax_cross_entropy_with_logits(
+        tf_f, tf_l)
+    self.assertAllClose([0, 0], tf_loss)
+    self.assertAllCloseAccordingToType([[0], [0]], tf_gradient)
+
   @test_util.run_deprecated_v1
   def testNotMatrix(self):
     with self.cached_session():

@@ -143,15 +143,22 @@ class MatrixSolveOpGpu : public AsyncOpKernel {
                       done);
     OP_REQUIRES_ASYNC(
         context, input.dim_size(ndims - 2) == n,
-        errors::InvalidArgument("Input matrices must be squares, got",
+        errors::InvalidArgument("Input matrices must be squares, got ",
                                 input.dim_size(ndims - 2), " != ", n),
         done);
     OP_REQUIRES_ASYNC(context, rhs.dim_size(ndims - 2) == n,
                       errors::InvalidArgument(
                           "Input matrix and right-hand side must have the "
-                          "same number of rows, got",
+                          "same number of rows, got ",
                           n, " != ", rhs.dim_size(ndims - 2)),
                       done);
+    for (int dim = 0; dim < ndims - 2; dim++) {
+      OP_REQUIRES_ASYNC(
+          context, input.dim_size(dim) == rhs.dim_size(dim),
+          errors::InvalidArgument(
+              "All input tensors must have the same outer dimensions."),
+          done);
+    }
 
     // Allocate output.
     Tensor* output;

@@ -112,6 +112,12 @@ def testWrongDimensions(self):
     with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
       self.evaluate(linalg_ops.matrix_solve(matrix, rhs))
 
+    # The matrix and right-hand side should have the same batch dimensions
+    matrix = np.random.normal(size=(2, 6, 2, 2))
+    rhs = np.random.normal(size=(2, 3, 2, 2))
+    with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
+      self.evaluate(linalg_ops.matrix_solve(matrix, rhs))
+
   def testNotInvertible(self):
     # The input should be invertible.
     with self.assertRaisesOpError("Input matrix is not invertible."):