 
     VisitDtype(out.dtype(), [&a, &indices, axis, &out, mode](auto pt) {
         using T = typename decltype(pt)::type;
-
-        // a and out are transposed as follows.
-        // a:       (Ni..., N,     Nj...) => (N,     Ni..., Nj...)
-        // out:     (Ni..., Nk..., Nj...) => (Nk..., Ni..., Nj...)
-        //
-        // indices is used as is.
-        // indices: (Nk...)
-
-        IndexableArray<const T> a_iarray{a};
-        Axes a_perm = MakeRollingPermutation(axis, axis + 1, a.ndim());
-        a_iarray.Permute(a_perm);
-        Shape a_shape = internal::TransposeShape(a.shape(), a_perm);
-        Indexer<> a_indexer{a_shape};
-
-        IndexableArray<T> out_iarray{out};
-        Axes out_perm = MakeRollingPermutation(axis, axis + indices.ndim(), out.ndim());
-        out_iarray.Permute(out_perm);
-        Shape out_shape = internal::TransposeShape(out.shape(), out_perm);
-        Indexer<> out_indexer{out_shape};
-
-        IndexableArray<const TIndex> indices_iarray{indices};
-        Indexer<> indices_indexer{indices.shape()};
-
-        // size of (Ni..., Nj...) part
-        TIndex common_total_size = gsl::narrow<TIndex>(a_indexer.total_size() / a_shape[0]);
-
-        TIndex axis_dim = gsl::narrow<TIndex>(a_shape[0]);
-
-        // TODO(niboshi): Calculate kMaxBlockSize per device
-        std::lock_guard<std::mutex> lock{*cuda_internal::g_mutex};
-        int64_t total_size = out_indexer.total_size();
-        static const int kWrapMaxBlockSize = CudaOccupancyMaxPotentialBlockSize(&TakeWrapCudaKernel<T, TIndex>).block_size;
-        static const int kClipMaxBlockSize = CudaOccupancyMaxPotentialBlockSize(&TakeClipCudaKernel<T, TIndex>).block_size;
-        int64_t grid_size;
-        int64_t block_size;
-        switch (mode) {
-            case IndexBoundsMode::kRaise:
-                throw BackendError{"Take with mode='raise' is not supported with CUDA backend"};
-            case IndexBoundsMode::kDefault:
-            case IndexBoundsMode::kWrap:
-                grid_size = (total_size + kWrapMaxBlockSize - 1) / kWrapMaxBlockSize;
-                block_size = std::min<TIndex>(total_size, kWrapMaxBlockSize);
-                TakeWrapCudaKernel<<<grid_size, block_size>>>(
-                        a_iarray, out_iarray, indices_iarray, a_indexer, out_indexer, indices_indexer, common_total_size, axis_dim);
-                break;
-            case IndexBoundsMode::kClip:
-                grid_size = (total_size + kClipMaxBlockSize - 1) / kClipMaxBlockSize;
-                block_size = std::min<TIndex>(total_size, kClipMaxBlockSize);
-                TakeClipCudaKernel<<<grid_size, block_size>>>(
-                        a_iarray, out_iarray, indices_iarray, a_indexer, out_indexer, indices_indexer, common_total_size, axis_dim);
-                break;
-        }
+        int64_t num_indices = indices.GetTotalSize();
+        int64_t target_dim = a.shape()[axis];
+        int64_t num_iters = out.GetTotalSize();
+        int64_t right_dim = std::accumulate(a.shape().begin() + axis + 1, a.shape().end(), int64_t{1}, std::multiplies<>());
+        static const int k1DMaxBlockSize = CudaOccupancyMaxPotentialBlockSize(&TakeCudaKernel<T, TIndex, 1>).block_size;
+        static const int kNDMaxBlockSize = CudaOccupancyMaxPotentialBlockSize(&TakeCudaKernel<T, TIndex, kDynamicNdim>).block_size;
+        if (a.IsContiguous() && indices.IsContiguous()) {
+            Array af = a.Ravel();
+            Array indicesf = indices.Ravel();
+            Array outf = out.Ravel();
+            IndexableArray<const T, 1> a_iarray{af};
+            IndexableArray<const TIndex, 1> indices_iarray{indicesf};
+            IndexableArray<T, 1> out_iarray{outf};
+            Indexer<1> a_indexer{af.shape()};
+            Indexer<1> indices_indexer{indicesf.shape()};
+            Indexer<1> out_indexer{outf.shape()};
+
+            TakeCudaKernel<T, TIndex, 1><<<(num_iters + k1DMaxBlockSize - 1) / k1DMaxBlockSize, k1DMaxBlockSize>>>(
+                    a_iarray,
+                    indices_iarray,
+                    out_iarray,
+                    a_indexer,
+                    indices_indexer,
+                    out_indexer,
+                    num_indices,
+                    target_dim,
+                    right_dim,
+                    num_iters,
+                    mode);
+        } else {
+            IndexableArray<const T> a_iarray{a};
+            IndexableArray<const TIndex> indices_iarray{indices};
+            IndexableArray<T> out_iarray{out};
+            Indexer<> a_indexer{a.shape()};
+            Indexer<> indices_indexer{indices.shape()};
+            Indexer<> out_indexer{out.shape()};
+
+            TakeCudaKernel<T, TIndex, kDynamicNdim><<<(num_iters + kNDMaxBlockSize - 1) / kNDMaxBlockSize, kNDMaxBlockSize>>>(
+                    a_iarray,
+                    indices_iarray,
+                    out_iarray,
+                    a_indexer,
+                    indices_indexer,
+                    out_indexer,
+                    num_indices,
+                    target_dim,
+                    right_dim,
+                    num_iters,
+                    mode);
+        }
+        return;
     });
 }
 

    return permutation;
}

template <typename T, typename TIndex>
__global__ void TakeWrapCudaKernel(
        IndexableArray<const T> a_iarray,
        IndexableArray<T> out_iarray,
        IndexableArray<const TIndex> indices_iarray,
        Indexer<> a_indexer,
        Indexer<> out_indexer,
        Indexer<> indices_indexer,
        TIndex common_total_size,
        TIndex axis_dim) {
    static_assert(std::is_same<TIndex, int64_t>::value || std::is_same<TIndex, int32_t>::value, "");
    for (auto it = out_indexer.It(blockIdx.x * blockDim.x + threadIdx.x, blockDim.x * gridDim.x); it; ++it) {
        TIndex indices_pos = static_cast<TIndex>(it.raw_index()) / common_total_size;
        TIndex common_pos = static_cast<TIndex>(it.raw_index()) % common_total_size;

        TIndex index = indices_iarray[indices_indexer.It(indices_pos)];
        if (index < 0) {
            index = axis_dim - ((-index + axis_dim - 1) % axis_dim + 1);
        } else {
            index = index % axis_dim;
        }
        CHAINERX_ASSERT(0 <= index);
        CHAINERX_ASSERT(index < axis_dim);

        out_iarray[it] = a_iarray[a_indexer.It(index * common_total_size + common_pos)];
    }
}

template <typename T, typename TIndex>
__global__ void TakeClipCudaKernel(
        IndexableArray<const T> a_iarray,
        IndexableArray<T> out_iarray,
        IndexableArray<const TIndex> indices_iarray,
        Indexer<> a_indexer,
        Indexer<> out_indexer,
        Indexer<> indices_indexer,
        TIndex common_total_size,
        TIndex axis_dim) {
    static_assert(std::is_same<TIndex, int64_t>::value || std::is_same<TIndex, int32_t>::value, "");
    for (auto it = out_indexer.It(blockIdx.x * blockDim.x + threadIdx.x, blockDim.x * gridDim.x); it; ++it) {
        TIndex indices_pos = static_cast<TIndex>(it.raw_index()) / common_total_size;
        TIndex common_pos = static_cast<TIndex>(it.raw_index()) % common_total_size;

        TIndex index = indices_iarray[indices_indexer.It(indices_pos)];
        index = max(TIndex{0}, min(index, axis_dim - 1));
        CHAINERX_ASSERT(0 <= index);
        CHAINERX_ASSERT(index < axis_dim);

        out_iarray[it] = a_iarray[a_indexer.It(index * common_total_size + common_pos)];
    }
}

template <typename T, typename TIndex>


template <typename TIndex>
void TakeImpl(Device& device, const Array& a, const Array& indices, int8_t axis, const Array& out, IndexBoundsMode mode) {
    if (mode == IndexBoundsMode::kRaise) {
        throw BackendError{"Take with mode='raise' is not supported with CUDA backend"};
    }
    static_assert(std::is_same<TIndex, int64_t>::value || std::is_same<TIndex, int32_t>::value, "");
    CHAINERX_ASSERT(
            (std::is_same<TIndex, int64_t>::value && indices.dtype() == Dtype::kInt64) ||
