 	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();

#include <ShapeUtils.h>
#include <numeric>
#include <PointersManager.h>

namespace nd4j    {
namespace ops     {

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

template<typename T>
__global__ void softMaxForVectorCudaGlobal(const void *vx, const Nd4jLong *xShapeInfo, void *vz, const Nd4jLong *zShapeInfo) {

	softMaxForVectorCuda<T>(vx, xShapeInfo, vz, zShapeInfo);
}

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

	softMaxForVectorCudaGlobal<T><<<1, MAX_NUM_THREADS / 4 , (MAX_NUM_THREADS / 4) * sizeof(T) + 512, *stream>>>(vx, xShapeInfo, vz, zShapeInfo);
}

///////////////////////////////////////////////////////////////////
template<typename T>
__global__ static void softMaxCuda(const void* vx, const Nd4jLong *xTadShapeInfo, const Nd4jLong *xOffsets,
                                         void* vz, const Nd4jLong *zTadShapeInfo, const Nd4jLong *zOffsets) {

    const auto x = reinterpret_cast<const T*>(vx);
          auto z = reinterpret_cast<T*>(vz);

    const auto* xTad = x + xOffsets[blockIdx.x];
          auto* zTad = z + zOffsets[blockIdx.x];

    softMaxForVectorCuda<T>(xTad, xTadShapeInfo, zTad, zTadShapeInfo);
}

///////////////////////////////////////////////////////////////////
template<typename T>
static void softMaxCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem, const cudaStream_t *stream,
                                const void* vx, const Nd4jLong *xTadShapeInfo, const Nd4jLong *xOffsets,
                                	  void* vz, const Nd4jLong *zTadShapeInfo, const Nd4jLong *zOffsets) {

    softMaxCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xTadShapeInfo, xOffsets, vz, zTadShapeInfo, zOffsets);
}


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