 
 /* Reduce one of the outer dimensions of a tensor
  *
- * For an n-d tensor (n <= 4) where the reduction is *not* along the innermost
- * dimension:
+ * For an n-d tensor where the reduction is *not* along the innermost dimension:
  *
- * - block.x and grid.x make up the innermost dimension;
- * - The reduced dimension is looped over inside a block; and
- * - grid.y and grid.z are the remaining two dimensions (if any).
- * - block.y and block.z are not used as we're limited to 512 or 1024 threads
- *   in the block.
- *
- * For sizes/strides, index 3 is the reduced dimension, while the remaining
- * indices are for the remaining dimensions with index 0 the innermost dimension.
+ * - blockIdx.x loops over the dimensions on the outside of the reduced dimension.
+ * - blockIdx.y and threadIdx.x loop over the dimensions on the inside of the
+ *   reduced dimension.
+ * - Each thread sequentially reduces one row of the reduced dimension.
  *
  * Reduction along the innermost dimension is handled in a separate kernel.
  */
 template<class UnaryFunction, class BinaryFunction>
 __global__ void THCudaTensor_kernel_transformReduceOuterDim(float *tgt, float *src_,
-        dim4 src_stride, dim4 tgt_stride, dim4 size,
+        unsigned num_orows, unsigned num_irows, unsigned row_size,
         UnaryFunction unary_op, float init, BinaryFunction binary_op)
 {
-  const size_t reduce = 3;
-
-  for(unsigned z = blockIdx.z; z < size[2] ; z += gridDim.z)
-  for(unsigned y = blockIdx.y; y < size[1] ; y += gridDim.y)
-  for(unsigned col = blockIdx.x * blockDim.x + threadIdx.x; col < size[0]; col += blockDim.x * gridDim.x) {
-    float *src = src_ + z * src_stride[2] + y * src_stride[1] + col;
-    float acc = init;
-    for(unsigned i=0; i < size[reduce]; i++) {
-      acc = binary_op(acc, unary_op(*src));
-      src += src_stride[reduce];
+
+  for (unsigned orow = blockIdx.x; orow < num_orows; orow += gridDim.x) {
+    for (unsigned irow = blockIdx.y * blockDim.x + threadIdx.x; irow < num_irows; irow += gridDim.y * blockDim.x) {
+      float *src = src_ + orow * row_size * num_irows + irow;
+      float acc = init;
+
+      for (unsigned col = 0; col < row_size; ++col) {
+        acc = binary_op(acc, unary_op(*src));
+        src += num_irows;
+      }
+      tgt[orow * num_irows + irow] = float(acc);
     }
-    tgt[z * tgt_stride[2] + y * tgt_stride[1] + col] = float(acc);
   }
 }
 

  return THCudaTensor_scanDim(state, self, src, dimension, 1.0f, thrust::multiplies<float>());
}

/* a set of reduction kernels that take in Binary ops on thrust pairs (of value, index)
   These are useful when you not only have to do a reduction, but you might have
   to preserve the location of contention (for example min/max operations)
 */
template<class BinaryFunction>
__global__ void THCudaTensor_kernel_transformReduceOuterDimIndex(float *tgt1, float *tgt2,
                                                             float *src_,
                                                             dim4 src_stride,
                                                             dim4 tgt1_stride,
                                                             dim4 tgt2_stride,
                                                             dim4 size,
                                                             thrust::pair<float,float> init,
                                                             BinaryFunction binary_op)
{
  const size_t reduce = 3;

  for(unsigned z = blockIdx.z; z < size[2] ; z += gridDim.z)
  for(unsigned y = blockIdx.y; y < size[1] ; y += gridDim.y)
  for(unsigned col = blockIdx.x * blockDim.x + threadIdx.x; col < size[0]; col += blockDim.x * gridDim.x) {
    float *src = src_ + z * src_stride[2] + y * src_stride[1] + col;
    thrust::pair<float,float> acc = init;
    for(unsigned i=0; i < size[reduce]; i++) {
      acc = binary_op(thrust::make_pair(*src, i+1), acc); // i+1 for 1-indexing
      src += src_stride[reduce];
    }
    tgt1[z * tgt1_stride[2] + y * tgt1_stride[1] + col] = acc.first;
    tgt2[z * tgt2_stride[2] + y * tgt2_stride[1] + col] = acc.second;
  }
}


        long dimension, UnaryFunction unary_op, float init, BinaryFunction binary_op)
{
  THArgCheck(dimension >= 0 && dimension < THCudaTensor_nDimension(state, src), 3, "dimension out of range");

  THLongStorage *dim = THCudaTensor_newSizeOf(state, src);
  THLongStorage_set(dim, dimension, 1);
