 		in  =  &in_flattened;
 		BF_ASSERT(in_flattened.ndim == out_flattened.ndim,
 		          BF_STATUS_INTERNAL_ERROR);
+		// TODO: Use streams to support multiple non-contiguous batch dims
+		//         (Like in linalg.cu)
 		BF_ASSERT(in_flattened.ndim == 3, BF_STATUS_UNSUPPORTED_SHAPE);
 		BF_ASSERT_EXCEPTION(in_flattened.shape[0] == out_flattened.shape[0],
 		                    BF_STATUS_INVALID_SHAPE);

		                    BF_STATUS_UNSUPPORTED_STRIDE);
		BF_ASSERT_EXCEPTION(out->strides[ndim-2] % out->strides[ndim-1] == 0,
		                    BF_STATUS_UNSUPPORTED_STRIDE);
		size_t istride =  in->strides[ndim-2] / in->strides[ndim-1];
		size_t ostride = out->strides[ndim-2] / in->strides[ndim-1];
		BF_ASSERT_EXCEPTION( in->strides[in->ndim-2] > 0, BF_STATUS_UNSUPPORTED_STRIDE);
		BF_ASSERT_EXCEPTION(out->strides[in->ndim-2] > 0, BF_STATUS_UNSUPPORTED_STRIDE);
		

				ostride_cur += reverse_time ? +1 : -1;
				obatchstride_cur = obatchstride;
			}
			launch_fdmt_exec_kernel(ntime, nrow, nbatch, (step==nstep-1), reverse_time,
			                        _d_step_delays  + step*_plan_stride,
			                        _d_step_srcrows + step*_plan_stride,
			                        d_ibuf, _buffer_stride, _batch_stride,
			                        d_obuf, ostride_cur, obatchstride_cur,
			                        _stream);
			std::swap(d_ibuf, d_obuf);
		}
		BF_CHECK_CUDA_EXCEPTION(cudaGetLastError(), BF_STATUS_INTERNAL_ERROR);
	}
	void set_stream(cudaStream_t stream) {
		_stream = stream;
