NPY Loader

This module implements the small, explicit .npy subset that TorchLean's native training examples use:

• NumPy format versions 1 and 2;
• little-endian float32 and float64 payloads (<f4, <f8);
• C-order arrays directly, and Fortran-order arrays converted to C-order at load time;
• deterministic conversion to a flat FloatArray in C order, or to one FloatArray per index along the leading axis.

The loader stays narrow. It is a runtime bridge for trusted experiment artifacts, not a general NumPy parser and not part of the formal tensor semantics. Tensor construction happens in NN.API.Data.Sources, which keeps this file-format boundary independent of model training.

Reference:

• NumPy .npy format documentation: https://numpy.org/doc/stable/reference/generated/numpy.lib.format.html

structure TorchLean.Data.IO.NpyData : Type

In-memory representation of a loaded .npy file in TorchLean's supported subset.

values is always flattened in C-order. If the source file declares fortran_order = True, we reorder the payload during parsing and store fortran := false in the returned value so downstream tensor loaders never have to reason about storage order.

• dtype : String

  Dtype string as stored in the header, for example "<f4" or "<f8".

• shape : Array Nat

  Logical array shape as stored in the header.

• fortran : Bool

  Whether the returned flat payload is still Fortran-ordered. This loader returns false.

• values : FloatArray

  Flattened numeric payload, converted to Lean Float values.

  This is a FloatArray rather than an Array Float because the payload is bulk numeric data: FloatArray stores the doubles unboxed, so a large file costs one machine word per element instead of one word plus one heap cell, and Tensor.from wraps it without copying.

structure TorchLean.Data.IO.NpyBlocks : Type

A .npy payload decoded as one FloatArray per index along the leading axis.

blocks.size is the leading dimension and every block holds the product of the trailing dimensions, so shape = #[d₀] ++ tail gives blocks.size = d₀ and blocks[i]!.size = tail.foldl (· * ·) 1.

• dtype : String

  Dtype string as stored in the header, for example "<f4" or "<f8".

• shape : Array Nat

  Logical array shape as stored in the header.

• blocks : Array FloatArray

  One decoded block per index along the leading axis, in order.

structure TorchLean.Data.IO.NpyHeader : Type

Validated NPY metadata, without reading or allocating the numeric payload.

Dimensions and the element count fit an addressable payload. Payload availability is checked separately by full or prefix readers; a valid header alone does not certify a complete file.

• dtype : String

  Supported dtype descriptor, either "<f4" or "<f8".

• shape : Array Nat

  Logical array shape from the header.

• fortran : Bool

  Whether the on-disk payload is Fortran-ordered.

• dataStart : Nat

  Byte offset where the numeric payload begins.

• elementBytes : Nat

  Number of bytes occupied by one on-disk element.

• elementCount : Nat

  Product of the validated dimensions; scalar arrays contain one element.

inductive TorchLean.Data.IO.Parallelism : Type

How a leading-axis block decode is spread across tasks.

• sequential : Parallelism

  Decode every block on the calling thread.

• tasks (n : Nat) : Parallelism

  Spread the blocks across at most n tasks.

• auto : Parallelism

  Spread the blocks across the platform's hardware concurrency.

def TorchLean.Data.IO.instReprParallelism.repr : Parallelism → Nat → Std.Format

@[instance_reducible]

instance TorchLean.Data.IO.instReprParallelism : Repr Parallelism

@[instance_reducible]

instance TorchLean.Data.IO.instBEqParallelism : BEq Parallelism

def TorchLean.Data.IO.instBEqParallelism.beq : Parallelism → Parallelism → Bool

def TorchLean.Data.IO.Internal.prefixProducts (shape : Array Nat) : Array Nat

Prefix products of a shape array.

For a shape [d₀, d₁, d₂], this returns [1, d₀, d₀*d₁], which are exactly the Fortran-order strides. We use these strides to convert Fortran storage into TorchLean's ordinary C-order flattening convention.

def TorchLean.Data.IO.Internal.idxFortranOfCIdx (shape : Array Nat) (idxC : Nat) : Nat

Convert a linear C-order index to the corresponding linear Fortran-order index.

Both indices describe the same multi-dimensional coordinate. The difference is only how the coordinate is flattened into a one-dimensional payload.

def TorchLean.Data.IO.Internal.reorderFortranToC (tag : String) (shape : Array Nat) (raw : FloatArray) : Except String FloatArray

Reorder a Fortran-ordered flat payload into C-order, rejecting an inconsistent payload.

The output stays unboxed, and an out-of-range index is an error rather than a 0.0 fill, so a malformed file is reported instead of silently reshaped.

def TorchLean.Data.IO.Internal.byteAt? (bs : ByteArray) (i : Nat) : Option UInt8

Safe ByteArray indexing.

def TorchLean.Data.IO.Internal.readLittleEndian (bytes : ByteArray) (offset byteCount : Nat) : Option Nat

Read byteCount little-endian bytes as a natural number, or none on truncated input.

def TorchLean.Data.IO.Internal.parseShapeValue (s : String) : Option (Array Nat)

Parse a shape tuple like (3, 4) or (3,) from a NumPy header fragment.

def TorchLean.Data.IO.Internal.parseHeader (tag hdr : String) : Except String (String × Bool × Array Nat)

Parse the NumPy header dictionary.

We only need three standard fields: descr, fortran_order, and shape. The header format is a Python-literal dictionary padded to an alignment boundary, so this parser stays field-oriented rather than trying to become a full Python parser.

def TorchLean.Data.IO.Internal.maxNpyHeaderBytes : Nat

Bound header allocation before interpreting an untrusted version-2 header length.

def TorchLean.Data.IO.Internal.npyHeaderLayout (tag : String) (bs : ByteArray) : Except String (Nat × Nat)

Validate the magic and version, returning the header offset and its bounded byte length.

def TorchLean.Data.IO.Internal.npyElementBytes (tag descr : String) : Except String Nat

Byte width for the dtypes supported by TorchLean's NPY loader.

def TorchLean.Data.IO.Internal.npyElementCount (tag : String) (shape : Array Nat) : Except String Nat

Bound each dimension and the shape product before allocating a decoded payload.

def TorchLean.Data.IO.Internal.parseNpyHeaderMeta (tag : String) (bs : ByteArray) : Except String NpyHeader

Parse and validate the metadata shared by in-memory and file-backed NPY readers.

Payload decoding

byteAt? and readLittleEndian are the right shape at the file boundary, where a malformed header has to be reported rather than guessed at. Inside a payload loop they are the wrong shape: an Option per byte and an Except per element each cost a heap cell, so decoding an array that way pays about ten allocations per element and runs an order of magnitude below the rate at which the same loop can fill a FloatArray.

The functions below decode a payload whose byte range and dtype the caller has already checked. The dtype is matched once per file rather than once per element, the bytes are read without an intermediate Option, and the result is unboxed.

The decoders reinterpret stored bits, with the one exception the runtime imposes: Float.ofBits and Float32.ofBits return the canonical quiet NaN for every NaN pattern, so a NaN decodes as a NaN but without its sign or payload. Signed zeros, subnormals, infinities and every finite value decode to exactly their stored bits.

@[inline]

def TorchLean.Data.IO.Internal.uint64LE (bs : ByteArray) (o : Nat) : UInt64

Little-endian UInt64 at byte offset o. Out-of-range bytes read as 0; callers check the payload length first, so that fallback is unreachable for accepted files.

@[inline]

def TorchLean.Data.IO.Internal.uint32LE (bs : ByteArray) (o : Nat) : UInt32

Little-endian UInt32 at byte offset o, with the same precondition as uint64LE.

def TorchLean.Data.IO.Internal.decodeRangeF8 (bs : ByteArray) (dataStart lo hi : Nat) : FloatArray

Decode the half-open element range [lo, hi) of a <f8 payload.

def TorchLean.Data.IO.Internal.decodeRangeF4 (bs : ByteArray) (dataStart lo hi : Nat) : FloatArray

Decode the half-open element range [lo, hi) of a <f4 payload, widening to Float.

def TorchLean.Data.IO.Internal.rangeDecoder (tag dtype : String) (bs : ByteArray) (dataStart : Nat) : Except String (Nat → Nat → FloatArray)

The range decoder for a supported dtype, selected once per file.

Returning the decoder rather than the decoded payload is what keeps the dtype test out of the element loop, and it lets one selection serve many ranges.

def TorchLean.Data.IO.Internal.minElementsPerTask : Nat

Element count below which spreading a decode across tasks does not pay.

A task costs a scheduling round trip that only a reasonably large slice earns back, so this bounds the task count from above by totalElements / minElementsPerTask.

def TorchLean.Data.IO.Internal.resolveTaskCount (p : Parallelism) (nBlocks totalElements : Nat) : Nat

Resolve a Parallelism request against a payload into an actual task count.

Three bounds apply, and the smallest wins:

• the caller's request, or the platform's hardware concurrency for auto;
• the number of blocks. A block is the smallest piece this can hand back whole, so raising the task count past nBlocks would mean cutting a block in half and copying the halves back together afterwards, and that copy costs more than the extra task saves. The leading dimension is therefore a hard ceiling on this decode's parallelism;
• totalElements / minElementsPerTask, so small payloads stay on one thread.

def TorchLean.Data.IO.Internal.npyPrefixHeader (tag : String) (hdr : NpyHeader) (expectedShape : Array Nat) : Except String NpyHeader

Validate a C-order leading prefix and retain only its requested shape and element count.

def TorchLean.Data.IO.Internal.npyBlocksLayout (tag : String) (hdr : NpyHeader) : Except String (Nat × Nat)

Validate that a header describes C-order leading-axis blocks, returning the block count and the element count per block.

Fortran-order and scalar arrays are rejected here, before any payload is read: a Fortran-order file interleaves every leading-axis block across the payload, so neither the contiguity nor the independence of the blocks holds, and a scalar has no leading axis to cut along.

def TorchLean.Data.IO.Internal.decodeNpyPayload (tag : String) (hdr : NpyHeader) (bs : ByteArray) : Except String NpyData

Decode a validated full or prefix payload, checking its byte range before reserving values.

def TorchLean.Data.IO.Internal.decodeNpyBlocks (tag : String) (hdr : NpyHeader) (bs : ByteArray) (parallelism : Parallelism) : Except String NpyBlocks

Decode a validated C-order payload as one block per leading-axis index, checking its byte range before reserving values.

For a C-order array of shape (d₀, d₁, …, dₖ) each leading-axis block is physically contiguous, so the d₀ blocks are independent and therefore support either sequential or concurrent decoding. parallelism chooses how many tasks that uses; see resolveTaskCount for the bounds that apply to the request. d₀ is a ceiling on it, so a file with a small leading axis is decoded with at most that many tasks however many cores are free.

def TorchLean.Data.IO.Internal.readNpyBytes (handle : IO.FS.Handle) (count : Nat) : IO (Except String ByteArray)

Read exactly the requested bytes using bounded allocations, rejecting an early end of file.

The buffer grows only with bytes actually read. An untrusted element count never becomes the capacity of a read buffer before the corresponding file contents have been observed.

partial def TorchLean.Data.IO.Internal.readNpyBytes.loop (handle : IO.FS.Handle) (remaining : Nat) (bytes : ByteArray) : IO (Except String ByteArray)

def TorchLean.Data.IO.Internal.readNpyHeaderFromHandle (handle : IO.FS.Handle) : IO (Except String NpyHeader)

Read only the bounded NPY preamble and header, leaving the handle at the numeric payload.

def TorchLean.Data.IO.Internal.readNpyPayload (handle : IO.FS.Handle) (hdr : NpyHeader) : IO (Except String NpyData)

Read and decode the validated payload at the handle's current position.

def TorchLean.Data.IO.parseNpy (tag : String) (bs : ByteArray) : Except String NpyData

Parse the bytes of a .npy file into NpyData.

The parser rejects unsupported dtypes, missing or malformed required header fields, and truncated payloads. Header parsing recognizes the three required fields; it is not a general Python parser. That makes loader failures explicit at the trust boundary instead of silently producing tensors with the wrong shape or partial data.

def TorchLean.Data.IO.parseNpyPrefix (tag : String) (expectedShape : Array Nat) (bs : ByteArray) : Except String NpyData

Parse only the requested leading rows of a C-order .npy array.

This supports large exported tensors kept on disk while a run uses only the first n rows. The rank and trailing dimensions must match exactly; only the leading axis may be larger than requested.

The implementation shares header and dtype parsing with parseNpy, then decodes only the requested prefix. This avoids building a full payload when a command asks for a small leading slice of a real image or sequence dataset. Only the requested range must be present; unused trailing rows are not checked for truncation.

Why C-order only? In row-major NPY files, the first n rows are physically contiguous, so the prefix is exactly the first n * trailingSize elements. In Fortran-order files the same logical prefix is interleaved across the payload, so prefix decoding would be unsound. Rather than silently returning bad rows, we reject Fortran-order prefix loading and ask callers to convert the array to C-order first.

def TorchLean.Data.IO.parseNpyBlocks (tag : String) (bs : ByteArray) (parallelism : Parallelism := Parallelism.auto) : Except String NpyBlocks

Parse a C-order .npy file into one FloatArray per index along the leading axis.

Use this when the consumer wants the blocks apart anyway: one channel, one band, one batch element at a time. The blocks are decoded concurrently under the bounds resolveTaskCount describes, and the result does not depend on the task count.

Use parseNpy when the consumer wants a single flat payload instead. That path decodes on one thread by design. Decoding it concurrently would mean filling several separate arrays and then copying all of them into one contiguous result, and that final copy walks every element a second time: it costs several times the concurrent decode it was supposed to speed up, and it holds the pieces and the result in memory simultaneously. Splitting the work is only worth it when the pieces are what the caller wanted.

Fortran-order and scalar arrays are rejected, as npyBlocksLayout explains.

def TorchLean.Data.IO.readNpyHeader (path : System.FilePath) : IO (Except String NpyHeader)

Read validated metadata without reading the payload.

Supports version 1 and 2 headers up to one MiB. Dtype, storage-order syntax, dimensions, and shape products are validated; payload completeness belongs to readNpy or readNpyPrefix.

def TorchLean.Data.IO.readNpy (path : System.FilePath) : IO (Except String NpyData)

Read the full declared NPY payload, rejecting truncation and converting Fortran to C order.

def TorchLean.Data.IO.readNpyPrefix (path : System.FilePath) (expectedShape : Array Nat) : IO (Except String NpyData)

Read only the header and requested leading rows of a C-order .npy file.

Unrequested rows are neither read nor decoded. The requested range must be complete, and rank and trailing dimensions must match exactly. Reads use bounded chunks so a malformed claimed size cannot trigger a correspondingly large allocation before any bytes have been read.

def TorchLean.Data.IO.readNpyBlocks (path : System.FilePath) (parallelism : Parallelism := Parallelism.auto) : IO (Except String NpyBlocks)

Read a C-order .npy file as one FloatArray per index along the leading axis.

The header is validated for block loading before the payload is read, so a Fortran-order or scalar file is rejected without reading its payload. The payload itself is read with the same bounded chunks as readNpy, and truncation is rejected.