Implementation:Ollama Ollama MLXRunner Ops

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

Overview

Go bindings for core MLX array operations, providing methods on the Array type for mathematical, shape manipulation, and tensor operations.

Description

Each method wraps a corresponding MLX C function via CGo. Includes arithmetic (Add, Multiply, Divide, Power, Subtract), shape operations (Reshape, Transpose, ExpandDims, Flatten, Squeeze), reduction (Argmax, SumAxis, Logsumexp), matrix operations (Matmul, GatherMM, Addmm), element-wise (Abs, Negative, Sigmoid, Tanh, Softmax), and array manipulation (Concatenate, TakeAxis, PutAlongAxis).

Usage

The operational foundation of the MLX runner. Every model forward pass, attention computation, and sampling operation is built from these primitives.

Code Reference

Source Location

• Repository: Ollama
• File: x/mlxrunner/mlx/ops.go
• Lines: 1-256

Signature

func (t *Array) Abs() *Array
func (t *Array) Add(other *Array) *Array
func (t *Array) Matmul(other *Array) *Array
func (t *Array) Multiply(other *Array) *Array
func (t *Array) Divide(other *Array) *Array
func (t *Array) Subtract(other *Array) *Array
func (t *Array) Reshape(shape ...int) *Array
func (t *Array) Transpose(axes ...int) *Array
func (t *Array) ExpandDims(axis int) *Array
func (t *Array) Squeeze(axis int) *Array
func (t *Array) Concatenate(axis int, others ...*Array) *Array
func (t *Array) Sigmoid() *Array
func (t *Array) Softmax(axis int) *Array
func (t *Array) Argmax(axis int, keepDims bool) *Array

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
t	*Array	Yes	The array to operate on (receiver)
other	*Array	Varies	Second operand for binary operations
axis	int	Varies	Axis for shape/reduction operations

Outputs

Name	Type	Description
result	*Array	New array with the operation result

Usage Examples

a := mlx.FromValues([]float32{1, 2, 3, 4}, 2, 2)
b := mlx.FromValues([]float32{5, 6, 7, 8}, 2, 2)

sum := a.Add(b)
product := a.Matmul(b)
reshaped := a.Reshape(4)
transposed := a.Transpose(1, 0)

Related Pages

• Principle:Ollama_Ollama_MLXRunner_Architecture