Implementation:Ollama Ollama Imagegen MLX Go

Knowledge Sources	Ollama
Domains	Image Generation, Native Bindings
Last Updated	2025-02-15 00:00 GMT

Overview

Core Go bindings for the MLX framework, providing array operations, memory management, and GPU acceleration via CGo.

Description

The mlx.go file is the primary Go interface to Apple's MLX framework, wrapping the C API through CGo. It defines the Array type with deterministic memory management (arrays tracked in a pool, freed on Eval), Dtype constants for all supported data types, and hundreds of tensor operations (matmul, add, mul, conv2d, attention, RoPE, quantization, etc.). The memory model uses Keep() to mark persistent arrays (weights, caches) and automatic cleanup of intermediates during Eval(). Stream management enables dedicated GPU streams for generation. The file also provides Metal GPU utilities for memory tracking and device configuration.

Usage

Used by all imagegen model implementations as the foundation for tensor operations, GPU computation, and memory management.

Code Reference

Source Location

• Repository: Ollama
• File: x/imagegen/mlx/mlx.go
• Lines: 1-2273

Signature

type Dtype int

const (
	DtypeBool      Dtype = C.MLX_BOOL
	DtypeFloat16   Dtype = C.MLX_FLOAT16
	DtypeFloat32   Dtype = C.MLX_FLOAT32
	DtypeBFloat16  Dtype = C.MLX_BFLOAT16
	// ... more types
)

type Array struct {
	c     C.mlx_array
	freed bool
	kept  bool
}

func newArray(array C.mlx_array) *Array
func NewArray(data []float32, shape []int32) *Array
func Eval(arrays ...*Array)
func Keep(arrays ...*Array)
func Matmul(a, b *Array) *Array
func ScaledDotProductAttention(q, k, v *Array, scale float32, causal bool) *Array
func RoPE(x *Array, dims int, traditional bool, base, scale float32, offset int) *Array

Import

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

I/O Contract

Inputs

Name	Type	Required	Description
arrays	[]*Array	Varies	MLX arrays for operations
shape	[]int32	Varies	Tensor dimensions
dtype	Dtype	Varies	Data type for array creation

Outputs

Name	Type	Description
*Array	*Array	Result array from operations (tracked for cleanup)

Usage Examples

// Create arrays
x := mlx.NewArray(data, []int32{1, 128})
w := mlx.NewArray(weights, []int32{128, 256})
mlx.Keep(w) // persist weight across Eval cycles

// Compute
y := mlx.Matmul(x, w)
mlx.Eval(y) // evaluates and frees non-kept intermediates

// Read result
result := y.Data()

Related Pages

• Principle:Ollama_Ollama_ImageGeneration