Implementation:Ollama Ollama MLXRunner Slice

Knowledge Sources	Ollama
Domains	MLX Runtime, Tensor Operations
Last Updated	2025-02-15 00:00 GMT

Overview

Provides Python-style array slicing operations for MLX arrays, supporting start:stop:stride notation across all dimensions.

Description

Defines a slice type and Slice(args ...int) constructor that supports four slice modes: Slice() for full dimension, Slice(i) for single index, Slice(i, j) for range, and Slice(i, j, k) for range with stride. The makeSlices helper converts slices to C-compatible start/stop/stride arrays. Array.Slice extracts sub-arrays and Array.SliceUpdate performs in-place updates on array slices.

Usage

Used extensively in KV cache management and attention operations where specific sequence positions need to be read or updated.

Code Reference

Source Location

• Repository: Ollama
• File: x/mlxrunner/mlx/slice.go
• Lines: 1-86

Signature

type slice struct {
    args []int
}

func Slice(args ...int) slice
func makeSlices(dims []int, slices ...slice) (starts, stops, strides []C.int)
func (t *Array) Slice(slices ...slice) *Array
func (t *Array) SliceUpdate(other *Array, slices ...slice) *Array

Import

import "github.com/ollama/ollama/x/mlxrunner/mlx"

I/O Contract

Inputs

Name	Type	Required	Description
slices	...slice	Yes	One slice per dimension defining the extraction range

Outputs

Name	Type	Description
result	*Array	Sliced sub-array

Usage Examples

// Extract rows 0 to 5 of a 2D tensor
sub := tensor.Slice(mlx.Slice(0, 5), mlx.Slice())

// Get single element at position [1, 2]
elem := tensor.Slice(mlx.Slice(1), mlx.Slice(2))

// Update a slice in-place
updated := tensor.SliceUpdate(newValues, mlx.Slice(), mlx.Slice(3, 6), mlx.Slice())

Related Pages

• Principle:Ollama_Ollama_MLXRunner_Architecture