gather | Modular

gather[dtype: DType, indices_type: DType, //, *, axis: Int, target: StringSlice[StaticConstantOrigin] = StringSlice("cpu")](output: TileTensor[dtype, output.LayoutType, output.origin, address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size], input: TileTensor[dtype, input.LayoutType, input.origin, address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size], indices: TileTensor[indices_type, indices.LayoutType, indices.origin, address_space=indices.address_space, linear_idx_type=indices.linear_idx_type, element_size=indices.element_size], *, context: DeviceContext)

Gather operation as defined in https://github.com/onnx/onnx/blob/main/docs/Operators.md#Gather.

Note that this is NOT the same as the default PyTorch gather (which is equivalent to https://github.com/onnx/onnx/blob/main/docs/Operators.md#gatherelements).

gather[dtype: DType, indices_type: DType, //, *, axis: Int, target: StringSlice[StaticConstantOrigin] = StringSlice("cpu")](output: TileTensor[dtype, output.LayoutType, output.origin, address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size], input: TileTensor[dtype, input.LayoutType, input.origin, address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size], indices: TileTensor[indices_type, indices.LayoutType, indices.origin, address_space=indices.address_space, linear_idx_type=indices.linear_idx_type, element_size=indices.element_size], *, context: DeviceContextPtr = DeviceContextPtr())

Gather operation as defined in https://github.com/onnx/onnx/blob/main/docs/Operators.md#Gather.

Note that this is NOT the same as the default PyTorch gather (which is equivalent to https://github.com/onnx/onnx/blob/main/docs/Operators.md#gatherelements).

gather[*, dtype: DType, indices_type: DType, input_fn: def[width: Int, rank: Int](IndexList[rank]) capturing -> SIMD[dtype, width], indices_fn: def[width: Int, rank: Int](IndexList[rank]) capturing -> SIMD[indices_type, width], output_fn: def[width: Int, rank: Int](IndexList[rank], SIMD[dtype, width]) capturing -> None, prefetch_fn: OptionalReg[def[input_rank: Int, indices_rank: Int](IndexList[input_rank], IndexList[indices_rank]) capturing -> None] = None, target: StringSlice[StaticConstantOrigin] = StringSlice("cpu")](axis: Axis, input_shape: IndexList[input_shape.size, element_type=input_shape.element_type], indices_shape: IndexList[indices_shape.size, element_type=indices_shape.element_type], output_shape: IndexList[output_shape.size, element_type=output_shape.element_type], *, context: DeviceContext)

Gather operation as defined in https://github.com/onnx/onnx/blob/main/docs/Operators.md#Gather.

Note that this is NOT the same as the default PyTorch gather (which is equivalent to https://github.com/onnx/onnx/blob/main/docs/Operators.md#gatherelements).

gather[*, dtype: DType, indices_type: DType, input_fn: def[width: Int, rank: Int](IndexList[rank]) capturing -> SIMD[dtype, width], indices_fn: def[width: Int, rank: Int](IndexList[rank]) capturing -> SIMD[indices_type, width], output_fn: def[width: Int, rank: Int](IndexList[rank], SIMD[dtype, width]) capturing -> None, prefetch_fn: OptionalReg[def[input_rank: Int, indices_rank: Int](IndexList[input_rank], IndexList[indices_rank]) capturing -> None] = None, target: StringSlice[StaticConstantOrigin] = StringSlice("cpu")](axis: Axis, input_shape: IndexList[input_shape.size, element_type=input_shape.element_type], indices_shape: IndexList[indices_shape.size, element_type=indices_shape.element_type], output_shape: IndexList[output_shape.size, element_type=output_shape.element_type], *, context: DeviceContextPtr = DeviceContextPtr())

Gather operation as defined in https://github.com/onnx/onnx/blob/main/docs/Operators.md#Gather.

Note that this is NOT the same as the default PyTorch gather (which is equivalent to https://github.com/onnx/onnx/blob/main/docs/Operators.md#gatherelements).