Sign in

    majikxu/laya/tensor does not have a README file

    Arena

    pub type Arena

    A flat float32 parameter arena held outside the MoonBit heap.

    ModernBERT-large expands to roughly 1.4 GiB of float32 weights, which is not something you want the collector walking on every cycle, so the arena is a single calloc behind an external object with a finalizer.

    Arena::capacity

    fn Arena::capacity(self : Arena) -> Int64

    Number of float32 slots in the arena.

    Arena::get

    fn Arena::get(self : Arena, index : Int64) -> Float

    Read one slot. Out-of-range reads return 0.0; this exists for tests and diagnostics, not for the hot path — the kernels address the arena natively.

    Arena::new

    fn Arena::new(len : Int64) -> Arena raise
    LayaError

    Allocate an arena of len float32 slots, zero-initialised.

    MetalBuffer

    pub type MetalBuffer

    A float32 buffer in GPU-visible memory, addressed by element offset.

    MetalBuffer::download

    fn MetalBuffer::download(self : MetalBuffer, dst : FixedArray[Float], count~ : Int, offset? : Int) -> Unit

    Copy count floats out of the buffer. Only valid after commit.

    MetalBuffer::upload

    fn MetalBuffer::upload(self : MetalBuffer, src : FixedArray[Float], count~ : Int, offset? : Int) -> Unit

    Copy count floats into the buffer at element offset offset.

    MetalBuffer::upload_ints

    fn MetalBuffer::upload_ints(self : MetalBuffer, src : FixedArray[Int], count~ : Int, offset? : Int) -> Unit

    Copy count ints into the buffer, for index operands such as the option marker positions.

    MetalContext

    pub type MetalContext

    A Metal device, command queue and the compiled kernel library.

    Operations encode into a command buffer rather than running: begin opens one, the operations queue work into it, and commit submits it and waits. A whole forward pass is therefore one submission and one synchronisation, which is what makes the GPU path worth having — per-op round trips would cost more than they save.

    MetalContext::add

    fn MetalContext::add(self : MetalContext, a : MetalBuffer, b : MetalBuffer, count~ : Int) -> Unit

    a += b over count elements.

    MetalContext::add_broadcast

    fn MetalContext::add_broadcast(self : MetalContext, x : MetalBuffer, v : MetalBuffer, rows~ : Int, cols~ : Int, v_offset~ : Int64) -> Unit

    Add one row vector to every row of [rows, cols].

    MetalContext::begin

    fn MetalContext::begin(self : MetalContext) -> Unit

    Open a command buffer. Operations encoded after this run on commit.

    MetalContext::buffer

    Allocate count float32 slots of GPU-visible storage.

    MetalContext::commit

    Submit the command buffer and wait for it. Buffer contents are readable afterwards.

    MetalContext::create

    fn MetalContext::create() -> Result[MetalContext, String]

    Create a Metal context, or report why it is unavailable.

    Returns Err rather than raising on an unsupported platform, because falling back to the CPU backend is a normal outcome, not a failure.

    MetalContext::device_name

    fn MetalContext::device_name(self : MetalContext) -> String

    Name of the GPU this context is bound to.

    MetalContext::fill_bias

    fn MetalContext::fill_bias(self : MetalContext, y : MetalBuffer, bias : MetalBuffer, rows~ : Int, cols~ : Int, y_offset? : Int, bias_offset~ : Int64) -> Unit

    Broadcast a bias row across [rows, cols], so a following gemm with beta = 1 adds it. MPS has no bias term of its own.

    MetalContext::gather_rows

    fn MetalContext::gather_rows(self : MetalContext, src : MetalBuffer, indices : MetalBuffer, dst : MetalBuffer, rows~ : Int, cols~ : Int) -> Unit

    Gather whole rows of src by index into dst.

    MetalContext::geglu

    fn MetalContext::geglu(self : MetalContext, x : MetalBuffer, y : MetalBuffer, rows~ : Int, inner~ : Int) -> Unit

    GeGLU over a [rows, 2 · inner] buffer.

    MetalContext::gelu

    fn MetalContext::gelu(self : MetalContext, x : MetalBuffer, count~ : Int) -> Unit

    Exact GELU in place.

    MetalContext::gemm

    fn MetalContext::gemm(self : MetalContext, a : MetalBuffer, b : MetalBuffer, c : MetalBuffer, m~ : Int, n~ : Int, k~ : Int, lda~ : Int, ldb~ : Int, ldc~ : Int, a_offset? : Int64, b_offset? : Int64, c_offset? : Int64, a_stride? : Int64, b_stride? : Int64, c_stride? : Int64, batch? : Int, trans_b? : Bool, alpha? : Float, beta? : Float) -> Unit

    C = alpha · A · op(B) + beta · C, over batch matrices with uniform strides. batch = 1 covers every Linear; batching covers attention, where each head is one matrix.

    MetalContext::layer_norm

    fn MetalContext::layer_norm(self : MetalContext, x : MetalBuffer, y : MetalBuffer, gamma : MetalBuffer, beta : MetalBuffer, rows~ : Int, cols~ : Int, x_offset? : Int, y_offset? : Int, gamma_offset~ : Int64, beta_offset~ : Int64, eps~ : Float) -> Unit

    Layer normalisation over [rows, cols]. beta_offset = -1 means no bias, in which case beta is ignored and may be any buffer.

    MetalContext::relu

    fn MetalContext::relu(self : MetalContext, x : MetalBuffer, count~ : Int) -> Unit

    ReLU in place.

    MetalContext::rope

    fn MetalContext::rope(self : MetalContext, x : MetalBuffer, heads~ : Int, len~ : Int, dim~ : Int, base~ : Double) -> Unit

    Rotary embedding in place over a packed [heads, len, dim] buffer.

    MetalContext::softmax

    fn MetalContext::softmax(self : MetalContext, scores : MetalBuffer, rows~ : Int, len~ : Int, window~ : Int) -> Unit

    Row-wise softmax over attention scores, applying the sliding-window mask in the same pass. window < 0 means unrestricted attention.

    MetalContext::split_qkv

    fn MetalContext::split_qkv(self : MetalContext, qkv : MetalBuffer, q : MetalBuffer, k : MetalBuffer, v : MetalBuffer, heads~ : Int, len~ : Int, dim~ : Int) -> Unit

    Unpack [len, 3, heads, dim] into three [heads, len, dim] buffers.

    MetalContext::unpack_heads

    fn MetalContext::unpack_heads(self : MetalContext, context_buffer : MetalBuffer, packed : MetalBuffer, heads~ : Int, len~ : Int, dim~ : Int) -> Unit

    Fold [heads, len, dim] back into [len, heads · dim].

    MetalContext::wrap_arena

    Expose the float32 weight arena to the GPU without copying it.

    Apple silicon shares memory between CPU and GPU, and the arena is mapped at a page boundary precisely so it can be wrapped here — ModernBERT-large would otherwise need a second 1.4 GB copy. The arena is kept alive for as long as the buffer exists.

    Reader

    pub type Reader

    A read-only handle on a safetensors checkpoint.

    Tensors are read with pread rather than mapped. Mapping looks attractive for the token embedding — 393 MiB on the multilingual checkpoint, and only ever indexed — but it keeps the whole file resident once loading has touched it, and Darwin will not hand those pages back. Reading keeps the footprint at the float32 arena plus a staging buffer.

    Reader::data_start

    fn Reader::data_start(self : Reader) -> Int64

    Byte offset at which tensor payloads begin.

    Reader::gather_rows_f16

    fn Reader::gather_rows_f16(self : Reader, offset : Int64, rows : FixedArray[Int], dim : Int, vocab : Int, out : FixedArray[Float]) -> Unit raise
    LayaError

    Gather rows.length() rows of width dim out of a float16 matrix of vocab rows located at offset, writing them contiguously into out.

    Reader::header

    fn Reader::header(self : Reader) -> Bytes

    The raw JSON header of the safetensors file.

    Reader::open

    fn Reader::open(path : String) -> Reader raise
    LayaError

    Open a safetensors file and validate its 8-byte header length prefix.

    Reader::read_f16

    fn Reader::read_f16(self : Reader, offset : Int64, count : Int64, arena : Arena, dst : Int64) -> Unit raise
    LayaError

    Decode count float16 values at offset into arena starting at element index dst.

    Reader::read_f16_into

    fn Reader::read_f16_into(self : Reader, offset : Int64, out : FixedArray[Float]) -> Unit raise
    LayaError

    Decode float16 values at offset into a MoonBit array, filling it entirely.

    Reader::read_f32

    fn Reader::read_f32(self : Reader, offset : Int64, count : Int64, arena : Arena, dst : Int64) -> Unit raise
    LayaError

    Copy count float32 values at offset into arena starting at dst.

    Reader::size

    fn Reader::size(self : Reader) -> Int64

    Total file size in bytes.

    add_inplace

    fn add_inplace(a : FixedArray[Float], b : FixedArray[Float], len~ : Int) -> Unit

    a += b over the first len elements.

    geglu

    fn geglu(x : FixedArray[Float], out : FixedArray[Float], rows~ : Int, inner~ : Int) -> Unit

    GeGLU over a [rows, 2 · inner] buffer produced by ModernBERT's Wi.

    Wi emits the value and gate halves concatenated along the last axis; out[rows, inner] = gelu(value) · gate.

    gelu

    fn gelu(x : Float) -> Float

    Exact GELU: 0.5 · x · (1 + erf(x / √2)).

    gelu_inplace

    fn gelu_inplace(x : FixedArray[Float], len~ : Int) -> Unit

    Apply exact GELU in place.

    gemm

    fn gemm(a : FixedArray[Float], b : FixedArray[Float], c : FixedArray[Float], m~ : Int, n~ : Int, k~ : Int, a_offset? : Int, b_offset? : Int, c_offset? : Int, lda~ : Int, ldb~ : Int, ldc~ : Int, trans_b? : Bool, alpha? : Float, beta? : Float) -> Unit

    C = alpha · A · op(B) + beta · C over activation arrays, with explicit offsets and leading dimensions so a single head's slice of a packed [heads, len, dim] buffer can be multiplied without copying.

    layer_norm

    fn layer_norm(x : FixedArray[Float], out : FixedArray[Float], rows~ : Int, cols~ : Int, gamma : FixedArray[Float], beta : FixedArray[Float]?, eps~ : Float) -> Unit

    Layer normalisation over the last axis of a [rows, cols] buffer.

    Means and variances accumulate in Double — at cols = 1024 the float32 sums drift enough to be visible against the reference.

    linear

    fn linear(x : FixedArray[Float], arena : Arena, weight_offset : Int64, bias_offset : Int64, y : FixedArray[Float], m~ : Int, k~ : Int, n~ : Int) -> Unit

    y[m, n] = x[m, k] · W[n, k]ᵀ (+ bias), with W and bias living in the arena at the given element offsets. Pass bias_offset = -1 for no bias.

    The transpose matches the PyTorch Linear layout the checkpoints use and is handed to BLAS rather than materialised.

    relu_inplace

    fn relu_inplace(x : FixedArray[Float], len~ : Int) -> Unit

    Apply ReLU in place.

    The decision head's feed-forward uses ReLU, not GELU: it is a torch.nn.TransformerEncoderLayer, whose default activation is ReLU, while the encoder and scorer around it are GELU.

    rope_inplace

    fn rope_inplace(x : FixedArray[Float], heads~ : Int, len~ : Int, dim~ : Int, base~ : Double) -> Unit

    Rotary position embedding in the "split halves" (NeoX / Hugging Face) layout, applied in place to a packed [heads, len, dim] buffer.

    For position p and pair i, the angle is p · base^(-2i/dim), and the pair rotated is (i, i + dim/2) rather than (2i, 2i+1).

    softmax

    fn softmax(logits : ArrayView[Double]) -> Array[Double]

    Softmax of logits into a fresh Array[Double].

    The decision layer works in Double once the encoder is done: the calibrated probabilities, the entropy-based confidence and the rounding that follows all compare against float64 reference values.

    softmax_rows

    fn softmax_rows(x : FixedArray[Float], rows~ : Int, cols~ : Int, offset? : Int) -> Unit

    In-place softmax over [rows, cols], row by row, shifted by the row max.

    zeros

    fn zeros(len : Int) -> FixedArray[Float]

    Allocate a zeroed activation buffer.