scatter.h File Reference

#include "mlx/backend/metal/kernels/indexing.h"

Go to the source code of this file.

Functions
template<typename T, typename IdxT, typename Op, int NIDX, bool UPD_ROW_CONTIG, int NWORK, typename LocT>
METAL_FUNC void	scatter_impl (const device T *updates, device mlx_atomic< T > *out, const constant int *upd_shape, const constant int64_t *upd_strides, const constant size_t &upd_ndim, const constant size_t &upd_size, const constant int *out_shape, const constant int64_t *out_strides, const constant size_t &out_ndim, const constant int *axes, const constant size_t &idx_size, const thread Indices< IdxT, NIDX > &indices, uint2 gid)

Function Documentation

scatter_impl()

template<typename T, typename IdxT, typename Op, int NIDX, bool UPD_ROW_CONTIG, int NWORK, typename LocT>

METAL_FUNC void scatter_impl	(	const device T *	updates,
		device mlx_atomic< T > *	out,
		const constant int *	upd_shape,
		const constant int64_t *	upd_strides,
		const constant size_t &	upd_ndim,
		const constant size_t &	upd_size,
		const constant int *	out_shape,
		const constant int64_t *	out_strides,
		const constant size_t &	out_ndim,
		const constant int *	axes,
		const constant size_t &	idx_size,
		const thread Indices< IdxT, NIDX > &	indices,
		uint2	gid )