 	if(!input.isActualOnDeviceSide()) input.syncToDevice();
 	const int rank = input.rankOf();
 
+	PointersManager manager(context, "helpers::softmax");
+
 	if(input.isVector()) {
 
 		if(rank == 1 || input.sizeAt(dimension) != 1) {
-			BUILD_SINGLE_SELECTOR(input.dataType(), softMaxForVectorCudaLauncher, (context->getCudaStream(), input.getSpecialBuffer(), input.getSpecialShapeInfo(), output.getSpecialBuffer()), FLOAT_TYPES);
-			input.tickReadDevice();
+			NDArray::prepareSpecialUse({&output}, {&input});
+			BUILD_SINGLE_SELECTOR(input.dataType(), softMaxForVectorCudaLauncher, (context->getCudaStream(), input.getSpecialBuffer(), input.getSpecialShapeInfo(), output.getSpecialBuffer(), output.getSpecialShapeInfo()), FLOAT_TYPES);
+			NDArray::registerSpecialUse({&output}, {&input});
 		}
 		else
 			output = 1.;
 	}
 	else {
 
-		auto maxAlongDim = const_cast<NDArray&>(input).reduceAlongDims(reduce::Max, {dimension}, true);
-		(input - maxAlongDim).applyTransform(transform::Exp, &output); // output contains exponents temporarily
-		auto sumAlongDim = output.reduceAlongDims(reduce::Sum, {dimension}, true);
-		output /= sumAlongDim;
-		input.tickReadDevice();
-	}
-
-	PointersManager manager(context, "helpers::softmax");
+		auto packX = nd4j::ConstantTadHelper::getInstance()->tadForDimensions(input.getShapeInfo(), {dimension});
+        auto packZ = nd4j::ConstantTadHelper::getInstance()->tadForDimensions(output.getShapeInfo(), {dimension});
+
+        const int threadsPerBlock = MAX_NUM_THREADS / 4;
+        const int blocksPerGrid = packZ.numberOfTads();
+        const int sharedMem = input.sizeOfT() * threadsPerBlock + 512;
+
+        NDArray::prepareSpecialUse({&output}, {&input});
+    	BUILD_SINGLE_SELECTOR(input.dataType(), softMaxCudaLauncher, (blocksPerGrid, threadsPerBlock, sharedMem, context->getCudaStream(), input.getSpecialBuffer(), packX.specialShapeInfo(), packX.specialOffsets(), output.specialBuffer(), packZ.specialShapeInfo(), packZ.specialOffsets()), FLOAT_TYPES);
+    	NDArray::registerSpecialUse({&output}, {&input});
+
+		// auto maxAlongDim = const_cast<NDArray&>(input).reduceAlongDims(reduce::Max, {dimension}, true);
+		// (input - maxAlongDim).applyTransform(transform::Exp, &output); // output contains exponents temporarily
+		// auto sumAlongDim = output.reduceAlongDims(reduce::Sum, {dimension}, true);
+		// output /= sumAlongDim;
+		// input.tickReadDevice();
+	}
+
+
 	manager.synchronize();
 
 	output.tickWriteDevice();

	for (int i = 0; i < numOfIters; ++i) {
		const Nd4jLong elemIdx = i * blockDim.x + threadIdx.x;
		if(elemIdx >= len) continue;
		const Nd4jLong offset = shape::getIndexOffset(elemIdx, xzShapeInfo, len);
		z[offset] /= shmem[0];
	}
}

///////////////////////////////////////////////////////////////////
template <typename T>
linkage void softMaxForVectorCudaLauncher(const cudaStream_t* stream, const void *vx, const Nd4jLong *xzShapeInfo, void *vz) {

	softMaxForVectorCuda<T><<<1, MAX_NUM_THREADS, MAX_NUM_THREADS * sizeof(T) + 512, *stream>>>(vx, xzShapeInfo, vz);
}

//////////////////////////////////////////////////////////////////////////
void softmax(nd4j::LaunchContext * context, const NDArray& input, NDArray& output, const int dimension) {

	if (threadIdx.x == 0) {
		extern __shared__ char shared[];
		shmem = reinterpret_cast<T*>(shared);
		len = shape::length(xShapeInfo);
		numOfIters = (len + blockDim.x - 1) / blockDim.x;   // ceil (len / blockDim.x)
	}
	__syncthreads();
