Skip to content

usage.md

Latest commit

 

History

History
352 lines (255 loc) · 7.04 KB

usage.md

File metadata and controls

352 lines (255 loc) · 7.04 KB
Raw

Usage Guide

mlx-go provides Go bindings for Apple's MLX framework, enabling high-performance machine learning on Apple Silicon.

Import

import mx "github.com/guptaaryan16/mlx_go/core"

Array Basics

Creating Arrays

// From Go slices
arr := mx.FromSlice([]float32{1.0, 2.0, 3.0, 4.0})
mat := mx.FromSlice2D([][]float32{{1, 2}, {3, 4}})

// Scalar arrays
scalar := mx.NewArrayFloat32(3.14)
intVal := mx.NewArrayInt(42)

// Creation functions
zeros := mx.Zeros([]int32{3, 3}, mx.Float32)
ones := mx.Ones([]int32{2, 4}, mx.Float32)
eye := mx.Eye(3, 3, 0, mx.Float32)           // 3x3 identity matrix
identity := mx.Identity(4, mx.Float32)       // 4x4 identity
arange := mx.Arange(0, 10, 1, mx.Float32)    // [0, 1, 2, ..., 9]
linspace := mx.Linspace(0, 1, 5, mx.Float32) // [0, 0.25, 0.5, 0.75, 1]

Array Properties

arr := mx.FromSlice([]float32{1, 2, 3, 4, 5, 6})
arr = mx.Reshape(arr, []int32{2, 3})

shape := arr.Shape()      // []int32{2, 3}
ndim := arr.Ndim()        // 2
size := arr.Size()        // 6
dtype := arr.Dtype()      // mx.Float32
itemSize := arr.ItemSize() // 4 (bytes per element)

Evaluation

MLX uses lazy evaluation. Call Eval() to trigger computation:

result := mx.Add(a, b)
result.Eval()  // Computation happens here
result.PrintArray()

Data Types

mx.Float32   // 32-bit floating point
mx.Float16   // 16-bit floating point
mx.BFloat16  // Brain floating point
mx.Int32     // 32-bit integer
mx.Int64     // 64-bit integer
mx.Uint32    // Unsigned 32-bit integer
mx.Bool      // Boolean

Mathematical Operations

Element-wise Operations

// Arithmetic
c := mx.Add(a, b)
c := mx.Subtract(a, b)
c := mx.Multiply(a, b)
c := mx.Divide(a, b)

// Unary operations
abs := mx.Abs(a)
neg := mx.Negative(a)
sqr := mx.Square(a)
sqrt := mx.Sqrt(a)
exp := mx.Exp(a)
log := mx.Log(a)

Matrix Operations

// Matrix multiplication: (m×n) @ (n×p) = (m×p)
C := mx.MatMul(A, B)

// Matrix-matrix multiply-add: C = alpha*A@B + beta*C
result := mx.AddMM(C, A, B, alpha, beta)

// Transpose
aT := mx.Transpose(a)

// Reshape
reshaped := mx.Reshape(a, []int32{4, 4})

Reductions

s := mx.Sum(a)
m := mx.Mean(a)
p := mx.Prod(a)
mn := mx.Min(a)
mx := mx.Max(a)

// Along specific axis
sumAxis := mx.SumAxis(a, 0, false)  // Sum along axis 0

Comparisons

result := mx.Equal(a, b)
result := mx.NotEqual(a, b)
result := mx.Greater(a, b)
result := mx.GreaterEqual(a, b)
result := mx.Less(a, b)
result := mx.LessEqual(a, b)

Neural Network Operations

Activation Functions

relu := mx.ReLU(x)
sigmoid := mx.Sigmoid(x)
softmax := mx.Softmax(x, axis)
logSoftmax := mx.LogSoftmax(x, axis)
gelu := mx.GELU(x)
silu := mx.SiLU(x)  // Swish activation
leakyRelu := mx.LeakyReLU(x, 0.01)

Normalization

// RMS Normalization (used in LLaMA, etc.)
normed := mx.RMSNorm(x, &weight, eps)

// Layer Normalization
normed := mx.LayerNorm(x, &weight, &bias, eps)

Attention

// Scaled Dot-Product Attention
output := mx.ScaledDotProductAttention(queries, keys, values, scale, &mask)

// Causal attention (for autoregressive models)
output := mx.ScaledDotProductAttentionCausal(queries, keys, values, scale)

Position Embeddings

// Rotary Position Embedding (RoPE)
output := mx.RoPE(x, dims, traditional, base, scale, offset)

// RoPE with custom frequencies
output := mx.RoPEWithFreqs(x, dims, traditional, scale, offset, freqs)

Convolutions

// 1D Convolution
out := mx.Conv1d(input, weight, stride, padding, dilation, groups)

// 2D Convolution
out := mx.Conv2d(input, weight, stride0, stride1, padding0, padding1, dilation0, dilation1, groups)

Embedding

// Lookup embeddings by index
embedded := mx.Embedding(weightMatrix, indices)

Linear Algebra

// Norms
l2 := mx.NormL2(a, nil, false)
pNorm := mx.Norm(a, 2.0, nil, false)

// Decompositions
Q, R := mx.QR(a)
svdResult := mx.SVD(a, true)  // Returns [U, S, Vt]
P, L, U := mx.LU(a)
chol := mx.Cholesky(a, false) // Lower triangular

// Matrix operations
inv := mx.Inv(a)
pinv := mx.PseudoInverse(a)
x := mx.Solve(A, b)  // Solve Ax = b

// Eigenvalues
vals, vecs := mx.Eig(a)
vals := mx.EigVals(a)

// Cross product
cross := mx.Cross(a, b, axis)

Random Number Generation

// Uniform distribution [low, high)
u := mx.RandomUniform(0, 1, []int32{3, 3}, mx.Float32)

// Normal distribution with mean=0, std=1
n := mx.RandomNormal([]int32{3, 3}, 0, 1, mx.Float32)

// Truncated normal
lower := mx.NewArrayFloat32(-2.0)
upper := mx.NewArrayFloat32(2.0)
tn := mx.RandomTruncatedNormal(lower, upper, []int32{3, 3}, mx.Float32)

// Random integers [low, high)
ints := mx.RandomRandint(0, 10, []int32{5}, mx.Int32)

// Random permutation
perm := mx.RandomPermutation(arr, axis)

// Categorical sampling from logits
samples := mx.RandomCategorical(logits, axis)

// Other distributions
gumbel := mx.RandomGumbel([]int32{3, 3}, mx.Float32)
laplace := mx.RandomLaplace([]int32{3, 3}, 0, 1, mx.Float32)

Streams and Devices

// Default stream (GPU if available)
result := mx.Add(a, b)  // Uses default stream

// Explicit GPU stream
result := mx.Add(a, b, mx.WithStream(mx.GPU))

// CPU stream
result := mx.Add(a, b, mx.WithStream(mx.CPU))

Memory Management

Arrays are automatically garbage collected via Go's runtime. For explicit control:

arr := mx.FromSlice([]float32{1, 2, 3})
// ... use arr ...
arr.Free()  // Optional: explicitly free C memory

Automatic Differentiation

MLX Go supports automatic gradient computation via Grad and ValueAndGrad:

// Define a differentiable function
loss := func(inputs ...mx.Array) []mx.Array {
    pred := mx.Add(mx.MatMul(inputs[2], inputs[0]), inputs[1])
    return []mx.Array{mx.Mean(mx.Square(mx.Subtract(pred, inputs[3])))}
}

// Compute gradients w.r.t. W (arg 0) and b (arg 1)
valueGradFn := mx.ValueAndGrad(loss, 0, 1)
values, grads := valueGradFn(W, b, X, Y)

// JIT compile for performance
fast := mx.Compile(loss)
fast = mx.Compile(loss, mx.Shapeless()) // shape-agnostic

See Autograd & Compilation for full documentation.

Converting Back to Go

// To 1D slice
floats := mx.ToSlice2D[float32](arr)

// For 2D arrays, returns flattened data row-major
data := mx.ToSlice2D[float32](mat2d)

Complete Example

package main

import (
    "fmt"
    mx "github.com/guptaaryan16/mlx_go/core"
)

func main() {
    // Create data
    X := mx.RandomUniform(0, 1, []int32{100, 10}, mx.Float32)
    W := mx.RandomNormal([]int32{10, 5}, 0, 0.1, mx.Float32)
    b := mx.Zeros([]int32{5}, mx.Float32)
    
    // Forward pass: Y = ReLU(X @ W + b)
    linear := mx.Add(mx.MatMul(X, W), b)
    Y := mx.ReLU(linear)
    Y.Eval()
    
    fmt.Print("Output shape: ")
    fmt.Println(Y.Shape()) // [100, 5]
}

Next Steps