 void AddAtImpl(Device& device, const Array& a, const Array& indices, int8_t axis, const Array& b, const Array& out, IndexBoundsMode mode) {
     // TODO(niboshi): Current implementation only distributes output elements in respective threads. Summation on the indices is performed
     // serially in each thread. This implementation can be improved by distributing indices as well, possibly using atomicAdd.
-
     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) ||

    out[out_indexer.It(yi)] = a[a_indexer.It(xi)];
}

template <typename T, typename TIndex>
__global__ void AddAtWrapCudaKernel(
        IndexableArray<const T> a_iarray,
        IndexableArray<const T> b_iarray,
        IndexableArray<T> out_iarray,
        IndexableArray<const TIndex> indices_iarray,
        Indexer<> b_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 axis_pos = static_cast<TIndex>(it.raw_index()) / common_total_size;
        TIndex common_pos = static_cast<TIndex>(it.raw_index()) % common_total_size;

        cuda_internal::DataType<T> out_value = cuda_internal::StorageToDataType<const T>(a_iarray[it]);

        for (auto it_indices = indices_indexer.It(0); it_indices; ++it_indices) {
            TIndex index = indices_iarray[it_indices];

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

            if (index == axis_pos) {
                out_value += cuda_internal::StorageToDataType<const T>(
                        b_iarray[b_indexer.It(it_indices.raw_index() * common_total_size + common_pos)]);
            }
        }

        out_iarray[it] = cuda_internal::DataToStorageType<T>(out_value);
    }
}

template <typename T, typename TIndex>
__global__ void AddAtClipCudaKernel(
        IndexableArray<const T> a_iarray,
        IndexableArray<const T> b_iarray,
        IndexableArray<T> out_iarray,
        IndexableArray<const TIndex> indices_iarray,
        Indexer<> b_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 axis_pos = static_cast<TIndex>(it.raw_index()) / common_total_size;
        TIndex common_pos = static_cast<TIndex>(it.raw_index()) % common_total_size;

        cuda_internal::DataType<T> out_value = cuda_internal::StorageToDataType<const T>(a_iarray[it]);

        for (auto it_indices = indices_indexer.It(0); it_indices; ++it_indices) {
            TIndex index = indices_iarray[it_indices];

            index = max(TIndex{0}, min(index, axis_dim - 1));
            CHAINERX_ASSERT(0 <= index);
            CHAINERX_ASSERT(index < axis_dim);

            if (index == axis_pos) {
                out_value += cuda_internal::StorageToDataType<const T>(
                        b_iarray[b_indexer.It(it_indices.raw_index() * common_total_size + common_pos)]);
            }
        }

        out_iarray[it] = cuda_internal::DataToStorageType<T>(out_value);
    }
}

Array CollapseNdArray(const Array& a) {


    VisitDtype(out.dtype(), [&a, &indices, axis, &b, &out, mode](auto pt) {
        using T = typename decltype(pt)::type;
        int64_t target_dim = a.shape()[axis];
        int64_t b_dim = b.shape()[axis];
        int64_t right_dim = 1;
        int64_t b_right_dim = 1;
        uint8_t uaxis = axis;
        int64_t num_iters = a.GetTotalSize();
        for (size_t i = axis; i < a.shape().size(); ++i) {
            if (i > uaxis) {
                right_dim *= a.shape()[i];
            }
        }
        for (size_t i = axis; i < b.shape().size(); ++i) {
            if (i > uaxis) {
                b_right_dim *= b.shape()[i];
            }
        }
        static const int k1DMaxBlockSize = CudaOccupancyMaxPotentialBlockSize(&AddAtCudaKernel<T, TIndex, 1>).block_size;
        static const int kNDMaxBlockSize = CudaOccupancyMaxPotentialBlockSize(&AddAtCudaKernel<T, TIndex, kDynamicNdim>).block_size;
        Array af = CollapseNdArray(a);
        Array bf = CollapseNdArray(b);
        Array indicesf = CollapseNdArray(indices);
        Array outf = CollapseNdArray(out);
        // To get the maximum performance we need to explicitly instantiate the 1D
        // version of the indexable array. If only one of the 3 arrays has more than one dim
        // the performance is greatly affected. In such case we just treat all the arrays equally
        // to simplify the code logic
        if (af.ndim() == 1 && bf.ndim() == 1 && indicesf.ndim() == 1 && outf.ndim() == 1) {
            IndexableArray<const T, 1> a_iarray{af};
            IndexableArray<const T, 1> b_iarray{bf};
            IndexableArray<const TIndex, 1> indices_iarray{indicesf};
            IndexableArray<T, 1> out_iarray{outf};
            Indexer<1> b_indexer{bf.shape()};
            Indexer<1> out_indexer{outf.shape()};
            Indexer<1> indices_indexer{indicesf.shape()};
            AddAtCudaKernel<T, TIndex, 1><<<(num_iters + k1DMaxBlockSize - 1) / k1DMaxBlockSize, k1DMaxBlockSize>>>(
                    a_iarray,
                    b_iarray,
                    out_iarray,
                    indices_iarray,
                    b_indexer,
                    out_indexer,
                    indices_indexer,
                    target_dim,
                    b_dim,
                    right_dim,
                    b_right_dim,
                    num_iters,
                    mode);
        } else {
            IndexableArray<const T> a_iarray{af};
            IndexableArray<const T> b_iarray{bf};
            IndexableArray<const TIndex> indices_iarray{indicesf};
            IndexableArray<T> out_iarray{outf};
            Indexer<> b_indexer{b.shape()};
            Indexer<> out_indexer{out.shape()};
            Indexer<> indices_indexer{indices.shape()};
            AddAtCudaKernel<T, TIndex, kDynamicNdim><<<(num_iters + kNDMaxBlockSize - 1) / kNDMaxBlockSize, kNDMaxBlockSize>>>(
                    a_iarray,
                    b_iarray,
                    out_iarray,
                    indices_iarray,
                    b_indexer,
                    out_indexer,
                    indices_indexer,
                    target_dim,
                    b_dim,
                    right_dim,
                    b_right_dim,
                    num_iters,
                    mode);
        }
    });
}
