Documentation
¶
Index ¶
- Constants
- Variables
- func ActiveMemory() int
- func AsyncEval(outputs ...*Array)
- func CUDAIsAvailable() bool
- func CacheMemory() int
- func CheckInit() error
- func ClearCache()
- func Compile1(name string, fn func(*Array) *Array, opts ...CompileOption) func(*Array) *Array
- func Compile2(name string, fn func(*Array, *Array) *Array, opts ...CompileOption) func(*Array, *Array) *Array
- func Compile3(name string, fn func(*Array, *Array, *Array) *Array, opts ...CompileOption) func(*Array, *Array, *Array) *Array
- func DisableCompile()
- func EnableCompile()
- func Eval(outputs ...*Array)
- func GPUIsAvailable() bool
- func GatedDelta(packed, ba, dtBias, aExp, state, mask *Array, captureAll bool) (y, nextState *Array, interior []*Array)
- func Load(path string) iter.Seq2[string, *Array]
- func LoadedLibraryPath() (string, error)
- func LogArrays()
- func Mamba2Scan(hidden, bState, cState, dt, state, a, d, dtBias, mask *Array, captureAll bool) (y, nextState *Array, interior []*Array)
- func MaxRecommendedWorkingSetSize() (int, error)
- func MetalIsAvailable() bool
- func PeakMemory() int
- func Pin(s ...*Array)
- func PrettyBytes(n int) fmt.Stringer
- func ResetPeakMemory()
- func SaveSafetensors(path string, arrays map[string]*Array) error
- func SaveSafetensorsWithMetadata(path string, arrays map[string]*Array, metadata map[string]string) error
- func SetDefaultDeviceGPU()
- func SetWiredLimit(limit int) (int, error)
- func Slice(args ...int) slice
- func Sweep()
- func Unpin(s ...*Array)
- func Version() string
- type Array
- func Add(a, b *Array) *Array
- func AddMM(c, a, b *Array, alpha, beta float32) *Array
- func AddScalar(a *Array, s float32) *Array
- func Argpartition(a *Array, kth int, axis int) *Array
- func Argsort(a *Array, axis int) *Array
- func Bernoulli(p *Array) *Array
- func BernoulliWithKey(p *Array, key *Array) *Array
- func Clamp(a *Array, minVal, maxVal float32) *Array
- func Clip(a, aMin, aMax *Array) *Array
- func Collect(v any) []*Array
- func Concatenate(arrays []*Array, axis int) *Array
- func Contiguous(a *Array, allowColMajor bool) *Array
- func Conv1d(x, weight *Array, bias *Array, stride, padding, dilation, groups int32) *Array
- func Conv2d(x, weight *Array, ...) *Array
- func Cos(a *Array) *Array
- func DepthwiseConv1d(x, weight *Array, bias *Array) *Array
- func DepthwiseConvSiLU(x, w, bias *Array, outLen int) *Array
- func Dequantize(w, scales, biases *Array, groupSize, bits int, mode string, globalScale *Array) *Array
- func Div(a, b *Array) *Array
- func DivScalar(a *Array, s float32) *Array
- func Erf(a *Array) *Array
- func Exp(a *Array) *Array
- func ExpandDims(a *Array, axis int) *Array
- func FastScaledDotProductAttention(q, k, v *Array, scale float32, mode string, mask *Array) *Array
- func Flatten(a *Array) *Array
- func FloorDivideScalar(a *Array, s int32) *Array
- func FromFP8(x *Array, dtype DType) *Array
- func FromValue[T scalarTypes](t T) *Array
- func FromValues[S ~[]E, E arrayTypes](s S, shape ...int) *Array
- func GLU(a *Array) *Array
- func GatherMM(a, b *Array, lhsIndices, rhsIndices *Array, sortedIndices bool) *Array
- func GatherQMM(x, w, scales *Array, biases, lhsIndices, rhsIndices *Array, transpose bool, ...) *Array
- func LayerNormFn(x, weight, bias *Array, eps float32) *Array
- func Log(a *Array) *Array
- func Logaddexp(a, b *Array) *Array
- func Matmul(a, b *Array) *Array
- func Maximum(a, b *Array) *Array
- func Mean(a *Array, axis int, keepDims bool) *Array
- func Minimum(a, b *Array) *Array
- func Mul(a, b *Array) *Array
- func MulScalar(a *Array, s float32) *Array
- func Neg(a *Array) *Array
- func New(name string) *Array
- func NewArrayInt32(data []int32, shape []int32) *Array
- func NewScalarArray(value float32) *Array
- func Pad(a *Array, axes []int, lowPad, highPad []int, padValue *Array, mode string) *Array
- func PadConstant(a *Array, axes []int, lowPad, highPad []int) *Array
- func Quantize(w *Array, groupSize, bits int, mode string) (weights, scales, biases *Array)
- func QuantizedMatmul(x, w, scales, biases *Array, transpose bool, groupSize, bits int, mode string) *Array
- func RMSNormFn(x, weight *Array, eps float32) *Array
- func RSqrt(a *Array) *Array
- func RandomKey(seed uint64) *Array
- func ReLU(a *Array) *Array
- func Reshape(a *Array, shape ...int32) *Array
- func RoPEWithBase(x *Array, dims int, traditional bool, base, scale float32, offsets *Array) *Array
- func RoPEWithFreqs(x *Array, dims int, traditional bool, base, scale float32, offsets *Array, ...) *Array
- func Sigmoid(a *Array) *Array
- func SigmoidRouter(gates, bias *Array) (origScores, negScores *Array)
- func Sin(a *Array) *Array
- func SliceStartStop(a *Array, start, stop []int32) *Array
- func SoftmaxAxis(a *Array, axis int, precise bool) *Array
- func Softplus(a *Array) *Array
- func Squeeze(a *Array, axis int) *Array
- func Stack(arrays []*Array, axis int) *Array
- func Sub(a, b *Array) *Array
- func Sum(a *Array, axis int, keepDims bool) *Array
- func Take(a *Array, indices *Array, axis int) *Array
- func TakeAlongAxis(a, indices *Array, axis int) *Array
- func Tile(a *Array, reps []int32) *Array
- func ToFP8(x *Array) *Array
- func Transpose(a *Array, axes ...int) *Array
- func Tri(n, m int32, k int) *Array
- func Where(condition, a, b *Array) *Array
- func Zeros(dtype DType, shape ...int) *Array
- func ZerosF32(shape []int32) *Array
- func (t *Array) Abs() *Array
- func (t *Array) Add(other *Array) *Array
- func (t *Array) Addmm(a, b *Array, alpha, beta float32) *Array
- func (t *Array) Argmax(axis int, keepDims bool) *Array
- func (t *Array) ArgpartitionAxis(kth int, axis int) *Array
- func (t *Array) ArgsortAxis(axis int) *Array
- func (t *Array) AsStrided(shape []int, strides []int, offset int) *Array
- func (t *Array) AsType(dtype DType) *Array
- func (t *Array) Categorical(axis int) *Array
- func (t *Array) CategoricalWithKey(axis int, key *Array) *Array
- func (t *Array) Clone() *Array
- func (t *Array) Concatenate(axis int, others ...*Array) *Array
- func (t *Array) Cumsum(axis int, reverse, inclusive bool) *Array
- func (t *Array) DType() DType
- func (t *Array) Dim(dim int) int
- func (t *Array) Dims() []int
- func (t *Array) Divide(other *Array) *Array
- func (t *Array) Equal(other *Array) *Array
- func (t *Array) ExpandDims(axis int) *Array
- func (t *Array) Flatten(startAxis, endAxis int) *Array
- func (t *Array) Float() float64
- func (t *Array) Floats() []float32
- func (t *Array) FloorDivide(other *Array) *Array
- func (t *Array) GatherMM(other, lhs, rhs *Array, sorted bool) *Array
- func (t *Array) Greater(other *Array) *Array
- func (t *Array) Int() int
- func (t *Array) Ints() []int
- func (t *Array) Less(other *Array) *Array
- func (t *Array) LessEqual(other *Array) *Array
- func (t *Array) LogValue() slog.Value
- func (t *Array) LogsumexpAxis(axis int, keepDims bool) *Array
- func (t *Array) Matmul(other *Array) *Array
- func (t *Array) MaxAxis(axis int, keepDims bool) *Array
- func (t *Array) Multiply(other *Array) *Array
- func (t *Array) Negative() *Array
- func (t *Array) NumBytes() int
- func (t *Array) NumDims() int
- func (t *Array) Power(exponent *Array) *Array
- func (t *Array) PutAlongAxis(indices, values *Array, axis int) *Array
- func (t *Array) Reshape(axes ...int) *Array
- func (t *Array) Save(name string) error
- func (t *Array) ScatterAddAxis(indices, values *Array, axis int) *Array
- func (t *Array) Set(other *Array)
- func (t *Array) Sigmoid() *Array
- func (t *Array) Size() int
- func (t *Array) Slice(slices ...slice) *Array
- func (t *Array) SliceUpdate(other *Array, slices ...slice) *Array
- func (t *Array) Sqrt() *Array
- func (t *Array) Squeeze(axis int) *Array
- func (t *Array) StackAxis(axis int, others ...*Array) *Array
- func (t *Array) String() string
- func (t *Array) Subtract(other *Array) *Array
- func (t *Array) SumAxis(axis int, keepDims bool) *Array
- func (t *Array) TakeAlongAxis(indices *Array, axis int) *Array
- func (t *Array) TakeAxis(indices *Array, axis int) *Array
- func (t *Array) Tanh() *Array
- func (t *Array) Transpose(axes ...int) *Array
- func (t *Array) Valid() bool
- type Byte
- type CompileFunc
- type CompileOption
- type DType
- type Device
- type Embedding
- type GibiByte
- type KibiByte
- type LayerNorm
- type Linear
- type MebiByte
- type Memory
- type RMSNorm
- type SafetensorsFile
- type Stream
- type TebiByte
Constants ¶
const End = math.MaxInt32
End is a sentinel value meaning "to the end of the dimension", equivalent to an omitted stop in Python (e.g. a[i:]).
Variables ¶
var GELU = Compile1("GELU", gelu, Shapeless())
GELU returns the exact erf formulation used by torch.nn.functional.gelu.
var GELUApprox = Compile1( "GELUApprox", func(x *Array) *Array { dt := x.DType() half := FromValue[float32](0.5).AsType(dt) coeff := FromValue(geluCoeff).AsType(dt) c := FromValue[float32](0.044715).AsType(dt) one := FromValue[float32](1.0).AsType(dt) x3 := x.Multiply(x).Multiply(x) inner := x.Add(c.Multiply(x3)) tanh := coeff.Multiply(inner).Tanh() return half.Multiply(x).Multiply(one.Add(tanh)) }, Shapeless(), )
GELUApprox returns 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) as a fused kernel.
var GeGLU = Compile2( "GeGLU", func(gate, up *Array) *Array { return GELUApprox(gate).Multiply(up) }, Shapeless(), )
GeGLU returns gelu_approx(gate) * up as a fused kernel. Matches mlx_lm's geglu, used by Gemma-family MLP and MoE paths.
var LogitSoftcap = Compile2( "LogitSoftcap", func(x, cap *Array) *Array { return x.Divide(cap).Tanh().Multiply(cap) }, Shapeless(), )
LogitSoftcap returns tanh(x / cap) * cap as a fused kernel. Matches mlx_lm's logit_softcap. cap must have the same dtype as x.
var ReLUSquared = Compile1( "ReLUSquared", func(x *Array) *Array { zero := FromValue[float32](0).AsType(x.DType()) x = Maximum(x, zero) return x.Multiply(x) }, Shapeless(), )
ReLUSquared returns relu(x)^2 as a fused kernel.
var SiLU = Compile1( "SiLU", func(a *Array) *Array { return a.Multiply(a.Sigmoid()) }, Shapeless(), )
SiLU returns a * sigmoid(a) as a fused kernel.
var SoftplusF32 = Compile1( "SoftplusF32", func(x *Array) *Array { dt := x.DType() zero := FromValue[float32](0) return Logaddexp(x.AsType(DTypeFloat32), zero).AsType(dt) }, Shapeless(), )
SoftplusF32 returns softplus(x) computed in float32 precision and cast back to x's original dtype, as a fused kernel. Matches the laguna attention output-gate formula: softplus(cast_f32(x)).cast(orig_dtype).
var SwiGLU = Compile2( "SwiGLU", func(gate, up *Array) *Array { return SiLU(gate).Multiply(up) }, Shapeless(), )
SwiGLU returns silu(gate) * up as a fused kernel.
Functions ¶
func ActiveMemory ¶
func ActiveMemory() int
func CUDAIsAvailable ¶ added in v0.32.6
func CUDAIsAvailable() bool
CUDAIsAvailable returns true if a CUDA GPU is available.
func CacheMemory ¶
func CacheMemory() int
func CheckInit ¶
func CheckInit() error
CheckInit returns any error that occurred during MLX dynamic library initialization.
func ClearCache ¶
func ClearCache()
func Compile2 ¶
func Compile2(name string, fn func(*Array, *Array) *Array, opts ...CompileOption) func(*Array, *Array) *Array
Compile2 compiles a binary function. See Compile.
func Compile3 ¶
func Compile3(name string, fn func(*Array, *Array, *Array) *Array, opts ...CompileOption) func(*Array, *Array, *Array) *Array
Compile3 compiles a ternary function. See Compile.
func DisableCompile ¶
func DisableCompile()
func EnableCompile ¶
func EnableCompile()
func GPUIsAvailable ¶
func GPUIsAvailable() bool
GPUIsAvailable returns true if a GPU device is available.
func GatedDelta ¶
func GatedDelta(packed, ba, dtBias, aExp, state, mask *Array, captureAll bool) (y, nextState *Array, interior []*Array)
GatedDelta runs the whole gated-delta step — q/k norms, decay gate, and the scan — in one launch over the activated causal-conv output. packed is [B, T, 2*Hk*Dk + Hv*Dv] with rows packed [q | k | v], ba is [B, T, 2*Hv] packed [beta | alpha] rows, and state is [B, Hv, Dv, Dk]. captureAll additionally emits every interior per-token state. Inputs that fit the one-launch kernels' contract run there; anything else runs the same step as graph ops.
When mask is non-nil, it must be a [B, T] bool tensor identifying real (true) vs. padded (false) positions. Padded rows are neutralized before the kernels' own preprocessing: zeroed conv rows make the q/k/v norms emit zeros, and -inf beta/alpha rows yield beta 0 and decay 1, so each padded position is an identity step with zero output, exactly as on the recurrence path.
func LoadedLibraryPath ¶ added in v0.30.0
LoadedLibraryPath returns the MLX dynamic library path selected by this package.
func Mamba2Scan ¶ added in v0.32.14
func Mamba2Scan(hidden, bState, cState, dt, state, a, d, dtBias, mask *Array, captureAll bool) (y, nextState *Array, interior []*Array)
Mamba2Scan runs the Mamba2 recurrent scan. Inputs must be float32 with shapes hidden [B, T, H, D], bState/cState [B, T, G, S] where H%G == 0, dt [B, T, H], state [B, H, D, S], and a/d/dtBias [H]. captureAll also returns the state after every token but the last.
mask, when non-nil, is a [B, T] bool marking real (true) vs padded positions; padded positions are identity steps.
func MaxRecommendedWorkingSetSize ¶ added in v0.32.4
MaxRecommendedWorkingSetSize returns the device's recommended upper bound for resident Metal allocations.
func MetalIsAvailable ¶
func MetalIsAvailable() bool
MetalIsAvailable returns true if a Metal GPU is available.
func PeakMemory ¶
func PeakMemory() int
func PrettyBytes ¶
func ResetPeakMemory ¶
func ResetPeakMemory()
func SaveSafetensors ¶
SaveSafetensors saves arrays to a safetensors file without metadata.
func SaveSafetensorsWithMetadata ¶
func SaveSafetensorsWithMetadata(path string, arrays map[string]*Array, metadata map[string]string) error
SaveSafetensorsWithMetadata saves arrays to a safetensors file with metadata.
func SetDefaultDeviceGPU ¶
func SetDefaultDeviceGPU()
SetDefaultDeviceGPU sets the default MLX device to GPU.
func SetWiredLimit ¶ added in v0.32.4
SetWiredLimit sets the maximum amount of Metal memory MLX keeps resident and returns the previous limit.
func Sweep ¶
func Sweep()
Sweep releases all unpinned arrays, primarily intermediate tensors. MLX will truly free them when there are no other references, including dependencies in the graph.
Types ¶
type Array ¶
type Array struct {
// contains filtered or unexported fields
}
func BernoulliWithKey ¶ added in v0.24.0
func Concatenate ¶
func Contiguous ¶
func Conv2d ¶
func Conv2d(x, weight *Array, strideH, strideW, padH, padW, dilationH, dilationW, groups int32) *Array
Conv2d performs 2D convolution: x [N,H,W,C_in], weight [C_out,kH,kW,C_in]. MLX uses NHWC layout.
func DepthwiseConv1d ¶
func DepthwiseConvSiLU ¶ added in v0.32.6
DepthwiseConvSiLU computes SiLU of a valid depthwise conv: x [B, T+K-1, C] and w [C, K] give [B, T, C], each output reading the K trailing input rows starting at its own index. bias, when non-nil, is [C]. Inputs that fit the fused kernel's contract run there; anything else runs the same computation as graph ops, bit for bit.
func Dequantize ¶
func ExpandDims ¶
func FastScaledDotProductAttention ¶ added in v0.22.1
func FloorDivideScalar ¶
func FromValues ¶
func GLU ¶
GLU applies Gated Linear Unit: splits x along last dim into two halves, returns first * sigmoid(second).
func LayerNormFn ¶
func NewArrayInt32 ¶
func NewScalarArray ¶
func Pad ¶
Pad pads array a along the given axes with specified low/high pad sizes. mode should be "constant", "edge", or "reflect".
func PadConstant ¶
PadConstant pads with zeros along the given axes.
func QuantizedMatmul ¶
func RoPEWithBase ¶
RoPEWithBase applies rotary position embeddings to x. offsets is an int32 array of shape [B] giving each batch row's starting position; the kernel applies positions offsets[b] + 0..T-1 per row.
func RoPEWithFreqs ¶
func RoPEWithFreqs(x *Array, dims int, traditional bool, base, scale float32, offsets *Array, freqs *Array) *Array
RoPEWithFreqs applies RoPE with optional custom frequencies. When freqs is non-nil, it is used instead of computing from base. Note: MLX takes reciprocal(freqs) internally to get inv_freq, so pass the actual frequencies (base^(2i/dim)), not the inverse frequencies.
func SigmoidRouter ¶
SigmoidRouter returns (sigmoid(gates), -(sigmoid(gates)+bias)) as a fused kernel — the DeepSeek-V2 / GLM-MoE aux-loss-free router head.
func SliceStartStop ¶
func TakeAlongAxis ¶
func (*Array) ArgsortAxis ¶
func (*Array) Categorical ¶
func (*Array) CategoricalWithKey ¶ added in v0.24.0
func (*Array) ExpandDims ¶
func (*Array) FloorDivide ¶
func (*Array) LogsumexpAxis ¶ added in v0.22.1
func (*Array) ScatterAddAxis ¶
func (*Array) SliceUpdate ¶
type CompileFunc ¶
CompileFunc is the signature of a function that can be compiled.
func Compile ¶
func Compile(name string, fn CompileFunc, opts ...CompileOption) CompileFunc
Compile returns a compiled version of fn. When called during another compile's trace, fn is inlined directly so outer compiles can fuse through inner ones.
Compiled functions must not have side effects outside of the function. Do not access data other than the arguments passed in (either Go data or MLX arrays) unless it is a constant.
type CompileOption ¶
type CompileOption func(*compileConfig)
CompileOption configures Compile behavior.
func Shapeless ¶
func Shapeless() CompileOption
Shapeless traces the function once against symbolic shapes so the compiled graph accepts any input shape afterwards. Without this option, MLX re-traces on each new (shape, dtype) combination and caches each specialization.
type DType ¶
type DType int
const ( DTypeBool DType = C.MLX_BOOL DTypeUint8 DType = C.MLX_UINT8 DTypeUint16 DType = C.MLX_UINT16 DTypeUint32 DType = C.MLX_UINT32 DTypeUint64 DType = C.MLX_UINT64 DTypeInt8 DType = C.MLX_INT8 DTypeInt16 DType = C.MLX_INT16 DTypeInt32 DType = C.MLX_INT32 DTypeInt64 DType = C.MLX_INT64 DTypeFloat16 DType = C.MLX_FLOAT16 DTypeFloat32 DType = C.MLX_FLOAT32 DTypeFloat64 DType = C.MLX_FLOAT64 DTypeBFloat16 DType = C.MLX_BFLOAT16 DTypeComplex64 DType = C.MLX_COMPLEX64 )
func (*DType) UnmarshalJSON ¶
type Device ¶
type Device struct {
// contains filtered or unexported fields
}
func DefaultDevice ¶
func DefaultDevice() Device
type Linear ¶
type SafetensorsFile ¶
type SafetensorsFile struct {
// contains filtered or unexported fields
}
SafetensorsFile represents a loaded safetensors file.
func LoadSafetensorsNative ¶
func LoadSafetensorsNative(path string) (*SafetensorsFile, error)
LoadSafetensorsNative loads a safetensors file using MLX's native loader.
func (*SafetensorsFile) Free ¶
func (s *SafetensorsFile) Free()
Free releases the loaded safetensors maps.
func (*SafetensorsFile) Get ¶
func (s *SafetensorsFile) Get(name string) *Array
Get retrieves a tensor by name.
func (*SafetensorsFile) GetMetadata ¶
func (s *SafetensorsFile) GetMetadata(key string) string
GetMetadata retrieves a metadata value by key.
type Stream ¶
type Stream struct {
// contains filtered or unexported fields
}
func DefaultStream ¶
func DefaultStream() Stream
Source Files
Directories