tensor_file.bend source
tensor_file.bend on the hub · documented module
# Raw little-endian FP32 files and NumPy .npy files through the official File effects.import Baseimport ./storage_buffer.bend as Buffersimport ./tensor.bend as Tensorsimport ./tensor_layout.bend as Layoutdef Tensor() -> Type: Tensors.Tensor<F32>def Read() -> Type: File & Result<&1,&1,U32 & String,List<&2,U32>>def Written() -> Type: File & Result<&1,&1,U32 & String,Unit># Values per File effect.def chunk() -> U32: 16384# An FP32 value and its 32 bits are the same word.def from_bits(word: U32) -> F32: U32{bits} = word F32{bits}def to_bits(value: F32) -> U32: F32{bits} = value U32{bits}# Reading: bytes gather into words, least significant first. The values of one# read are written into the buffer before the next read, so no list outgrows a read.type Words is Data: Words{values: List<&2,F32>,count: U32,word: U32,weight: U32}# The buffer of the tensor, the values found so far, and a word in progress.type Target is Type: Target{buffer: Buffers.Buffer<F32>,found: U32,word: U32,weight: U32}type Reading is Type: More{file: File,target: Target} End{file: File,target: Target,complete: Bool}type Outcome is Type: Outcome{target: Target,complete: Bool}def gather(full: Bool,values: List<&2,F32>,+count: U32,+word: U32,+weight: U32) -> Words: match full: case False{}: Words{values,count,word,(weight * 256 : U32)} case True{}: Words{from_bits(word) <> values,(count + 1 : U32),0,1}def push(+byte: U32,words: Words) -> Words: Words{values,+count,+word,+weight} = words gather(U32.is_eq(weight,16777216),values,count,(word + byte * weight : U32),weight)def trickle(bytes: List<&2,U32>,words: Words) -> Words: match bytes: case []: words case byte <> rest: trickle(rest,push(byte,words))# Four bytes at a time while the bytes are whole words; the values are newest first.def quads(bytes: List<&2,U32>,values: List<&2,F32>,+count: U32) -> Words: match bytes: case Con{first,Con{second,Con{third,Con{fourth,rest}}}}: quads(rest,from_bits((first + second * 256 + third * 65536 + fourth * 16777216 : U32)) <> values,(count + 1 : U32)) case rest: trickle(rest,Words{values,count,0,1})def words(whole: Bool,bytes: List<&2,U32>,+word: U32,+weight: U32) -> Words: match whole: case True{}: quads(bytes,[],0) case False{}: trickle(bytes,Words{[],0,word,weight})# Values beyond the count of the shape are counted and not written.def placed(read: Words,buffer: Buffers.Buffer<F32>,+found: U32,+count: U32) -> Target: Words{values,+added,+word,+weight} = read Target{Buffers.write_at(~F32,List.take(&2,F32,List.reverse(&2,F32,values),U32.to_nat((count - U32.min(found,count) : U32))),U32.min(found,count),buffer), (found + added : U32),word,weight}def fill(bytes: List<&2,U32>,target: Target,+count: U32) -> Target: Target{buffer,+found,+word,+weight} = target placed(words(U32.is_eq(weight,1),bytes,word,weight),buffer,found,count)def absorbed(bytes: List<&2,U32>,file: File,target: Target,+count: U32) -> Reading: match bytes: case []: End{file,target,True{}} case byte <> rest: More{file,fill(byte <> rest,target,count)}def received(result: Result<&1,&1,U32 & String,List<&2,U32>>,file: File,target: Target,+count: U32) -> Reading: match result: case Fail{error}: End{file,target,False{}} case Done{bytes}: absorbed(bytes,file,target,count)def taken(read: Read(),target: Target,+count: U32) -> Reading: (file,result) = read received(result,file,target,count)def closed(file: File,target: Target,complete: Bool) -> IO(Outcome): do IO<Outcome>: File.close(file) IO.pure(Outcome,Outcome{target,complete})# A file longer than its shape stops at the fuel and fails the count.def read_all(fuel: Nat,reading: Reading,+count: U32) -> IO(Outcome): match fuel reading: case 0n More{file,target}: closed(file,target,True{}) case 0n End{file,target,complete}: closed(file,target,complete) case 1n+rest End{file,target,complete}: closed(file,target,complete) case 1n+rest More{file,target}: do IO<Outcome>: read : Read() <- File.read_bytes(file,(chunk() * 4 : U32)) read_all(rest,taken(read,target,count),count)def fresh(+count: U32) -> Target: Target{Tensors.buffer(~F32,0.0,count),0,0,1}def fuel(+count: U32) -> Nat: U32.to_nat((count / chunk() + 2 : U32))def fault(operation: String,path: String,detail: String,shape: List<&2,U32>) -> Tensor(): Tensors.failed(~F32,operation,Tensors.FileFault{path,detail},shape,0.0,1)def sized(matches: Bool,operation: String,path: String,+shape: List<&2,U32>,buffer: Buffers.Buffer<F32>) -> Tensor(): match matches: case True{}: Tensors.from_storage(~F32,shape,buffer,0.0) case False{}: fault(operation,path,"does not hold the values of the shape",shape)def whole(aligned: Bool,operation: String,path: String,+shape: List<&2,U32>,buffer: Buffers.Buffer<F32>,+found: U32,+count: U32) -> Tensor(): match aligned: case False{}: fault(operation,path,"is not a whole number of FP32 values",shape) case True{}: sized(U32.is_eq(found,count),operation,path,shape,buffer)def aligned(operation: String,path: String,+shape: List<&2,U32>,target: Target,+count: U32) -> Tensor(): Target{buffer,+found,word,+weight} = target whole(U32.is_eq(weight,1),operation,path,shape,buffer,found,count)def checked(complete: Bool,operation: String,path: String,+shape: List<&2,U32>,target: Target,+count: U32) -> Tensor(): match complete: case False{}: fault(operation,path,"cannot be read",shape) case True{}: aligned(operation,path,shape,target,count)def loaded(operation: String,path: String,+shape: List<&2,U32>,outcome: Outcome,+count: U32) -> Tensor(): Outcome{target,complete} = outcome checked(complete,operation,path,shape,target,count)def opened(result: Result<&1,&1,U32 & String,File>,path: String,+shape: List<&2,U32>,+count: U32) -> IO(Tensor()): match result: case Fail{error}: IO.pure(Tensor(),fault("load",path,"cannot be opened",shape)) case Done{file}: do IO<Tensor()>: outcome : Outcome <- read_all(fuel(count),More{file,fresh(count)},count) IO.pure(Tensor(),loaded("load",path,shape,outcome,count))def counted(count: Maybe<&2,U32>,+path: String,+shape: List<&2,U32>) -> IO(Tensor()): match count: case None{}: IO.pure(Tensor(),Tensors.failed(~F32,"load",Tensors.InvalidShape{shape},shape,0.0,1)) case Some{+count}: do IO<Tensor()>: result : Result<&1,&1,U32 & String,File> <- File.open(path,"r") opened(result,path,shape,count)# The file holds the values of the shape in row-major order and nothing else.def load(+path: String,+shape: List<&2,U32>) -> IO(Tensor()): counted(Tensors.count_of(shape),path,shape)# The NumPy .npy header: magic, version, header length, then a text such as# {'descr': '<f4', 'fortran_order': False, 'shape': (16, 3, 6, 6), }type Header is Data: Header{shape: List<&2,U32>,data: List<&2,U32>}def codes(text: String) -> List<&2,U32>: match text: case SNil{}: [] case SCon{head,tail}: Char.to_u32(head) <> codes(tail)def magic() -> List<&2,U32>: [147,78,85,77,80,89]def starts(pattern: List<&2,U32>,bytes: List<&2,U32>) -> Bool: match pattern bytes: case Nil{} bytes: True{} case Con{first,rest} Nil{}: False{} case Con{first,rest} Con{byte,more}: U32.is_eq(first,byte) && starts(rest,more)def at(bytes: List<&2,U32>,index: Nat) -> U32: match bytes index: case Nil{} index: 0 case Con{byte,rest} 0n: byte case Con{byte,rest} 1n+before: at(rest,before)def contains(text: List<&2,U32>,+pattern: List<&2,U32>) -> Bool: match text: case Nil{}: False{} case Con{+byte,+rest}: starts(pattern,byte <> rest) || contains(rest,pattern)# The text after the first occurrence of the pattern.def after(text: List<&2,U32>,+pattern: List<&2,U32>) -> Maybe<&2,List<&2,U32>>: match text: case Nil{}: None{} case Con{+byte,+rest}: Bool.pick(Maybe<&2,List<&2,U32>>,starts(pattern,byte <> rest),Some{List.drop(&2,U32,byte <> rest,List.length(&2,U32,pattern))},after(rest,pattern))def pushed(digits: Bool,current: U32,shape: List<&2,U32>) -> List<&2,U32>: match digits: case True{}: current <> shape case False{}: shapedef closed_shape(valid: Bool,shape: List<&2,U32>) -> Maybe<&2,List<&2,U32>>: match valid: case True{}: Some{List.reverse(&2,U32,shape)} case False{}: None{}# Extents up to the closing parenthesis: digits, commas and spaces.def extents(text: List<&2,U32>,+current: U32,digits: Bool,valid: Bool,shape: List<&2,U32>) -> Maybe<&2,List<&2,U32>>: match text: case Nil{}: None{} case Con{+byte,rest}: match byte: case 41: closed_shape(valid,pushed(digits,current,shape)) case 44: extents(rest,0,False{},valid,pushed(digits,current,shape)) case 32: extents(rest,current,digits,valid,shape) case _: extents(rest,(current * 10 + byte - 48 : U32),True{},valid && U32.is_le(48,byte) && U32.is_le(byte,57),shape)def shaped(shape: Maybe<&2,List<&2,U32>>,data: List<&2,U32>) -> Maybe<&2,Header>: match shape: case None{}: None{} case Some{shape}: Some{Header{shape,data}}def typed(supported: Bool,text: Maybe<&2,List<&2,U32>>,data: List<&2,U32>) -> Maybe<&2,Header>: match supported text: case False{} text: None{} case True{} None{}: None{} case True{} Some{text}: shaped(extents(text,0,False{},True{},[]),data)def described(+text: List<&2,U32>,+length: Nat,data: List<&2,U32>) -> Maybe<&2,Header>: typed(Nat.is_eq(List.length(&2,U32,text),length) && contains(text,codes("'descr': '<f4'")) && contains(text,codes("'fortran_order': False")), after(text,codes("'shape': (")),data)def region(+bytes: List<&2,U32>,+length: Nat) -> Maybe<&2,Header>: described(List.take(&2,U32,bytes,length),length,List.drop(&2,U32,bytes,length))# Version 1 has a two-byte header length; later versions have four bytes.def versioned(major: U32,+bytes: List<&2,U32>) -> Maybe<&2,Header>: match major: case 1: region(List.drop(&2,U32,bytes,10n),U32.to_nat((at(bytes,8n) + at(bytes,9n) * 256 : U32))) case _: region(List.drop(&2,U32,bytes,12n),U32.to_nat((at(bytes,8n) + at(bytes,9n) * 256 + at(bytes,10n) * 65536 + at(bytes,11n) * 16777216 : U32)))def marked(valid: Bool,+bytes: List<&2,U32>) -> Maybe<&2,Header>: match valid: case False{}: None{} case True{}: versioned(at(bytes,6n),bytes)def header(+bytes: List<&2,U32>) -> Maybe<&2,Header>: marked(starts(magic(),bytes),bytes)def refused(file: File,tensor: Tensor()) -> IO(Tensor()): do IO<Tensor()>: File.close(file) IO.pure(Tensor(),tensor)def npy_counted(count: Maybe<&2,U32>,+shape: List<&2,U32>,data: List<&2,U32>,file: File,+path: String) -> IO(Tensor()): match count: case None{}: refused(file,Tensors.failed(~F32,"load_npy",Tensors.InvalidShape{shape},shape,0.0,1)) case Some{+count}: do IO<Tensor()>: outcome : Outcome <- read_all(fuel(count),More{file,fill(data,fresh(count),count)},count) IO.pure(Tensor(),loaded("load_npy",path,shape,outcome,count))def npy_parsed(parsed: Maybe<&2,Header>,file: File,+path: String) -> IO(Tensor()): match parsed: case None{}: refused(file,fault("load_npy",path,"is not a little-endian FP32 .npy array in C order",[])) case Some{Header{+shape,data}}: npy_counted(Tensors.count_of(shape),shape,data,file,path)def npy_received(result: Result<&1,&1,U32 & String,List<&2,U32>>,file: File,+path: String) -> IO(Tensor()): match result: case Fail{error}: refused(file,fault("load_npy",path,"cannot be read",[])) case Done{bytes}: npy_parsed(header(bytes),file,path)def npy_headed(read: Read(),+path: String) -> IO(Tensor()): (file,result) = read npy_received(result,file,path)def npy_opened(result: Result<&1,&1,U32 & String,File>,+path: String) -> IO(Tensor()): match result: case Fail{error}: IO.pure(Tensor(),fault("load_npy",path,"cannot be opened",[])) case Done{file}: do IO<Tensor()>: read : Read() <- File.read_bytes(file,(chunk() * 4 : U32)) npy_headed(read,path)# An FP32 array saved by NumPy, with its shape.def load_npy(+path: String) -> IO(Tensor()): do IO<Tensor()>: result : Result<&1,&1,U32 & String,File> <- File.open(path,"r") npy_opened(result,path)# Writingtype Writing is Type: Pending{file: File,storage: Array<F32>,offset: U32,remaining: Nat} Finished{file: File,complete: Bool}def bytes(+word: U32,rest: List<&2,U32>) -> List<&2,U32>: (word % 256 : U32) <> (word / 256 % 256 : U32) <> (word / 65536 % 256 : U32) <> (word / 16777216 : U32) <> restdef encode(count: Nat,values: List<&2,F32>) -> List<&2,U32>: match count values: case 0n values: [] case 1n+rest Nil{}: [] case 1n+rest Con{value,more}: bytes(to_bits(value),encode(rest,more))# The cursor retains the array; only the current chunk becomes a list.def read_chunk(count: Nat,+offset: U32,storage: Array<F32>,values: List<&2,F32>) -> Array<F32> & List<&2,F32>: match count: case 0n: (storage,List.reverse(&2,F32,values)) case 1n+rest: Buffers.bind_pair(Array<F32>,F32,Array<F32> & List<&2,F32>,Array.get(F32,storage,offset),storage => value => read_chunk(rest,(offset + 1 : U32),storage,value <> values))def pending(remaining: Nat,file: File,storage: Array<F32>,offset: U32) -> Writing: match remaining: case 0n: Finished{file,True{}} case 1n+rest: Pending{file,storage,offset,1n+rest}def continued(result: Result<&1,&1,U32 & String,Unit>,file: File,storage: Array<F32>,remaining: Nat,offset: U32) -> Writing: match result: case Fail{error}: Finished{file,False{}} case Done{unit}: pending(remaining,file,storage,offset)def after_write(written: Written(),storage: Array<F32>,remaining: Nat,offset: U32) -> Writing: (file,result) = written continued(result,file,storage,remaining,offset)def write_chunk(file: File,read: Array<F32> & List<&2,F32>,+count: Nat,remaining: Nat,offset: U32) -> IO(Writing): (storage,values) = read do IO<Writing>: written : Written() <- File.write_bytes(file,encode(count,values)) IO.pure(Writing,after_write(written,storage,Nat.sub(remaining,count),(offset + U32.from_nat(count) : U32)))def finished(file: File,complete: Bool) -> IO(Bool): do IO<Bool>: File.close(file) IO.pure(Bool,complete)def write_all(fuel: Nat,writing: Writing) -> IO(Bool): match fuel writing: case 0n Pending{file,storage,offset,remaining}: finished(file,False{}) case 0n Finished{file,complete}: finished(file,complete) case 1n+rest Finished{file,complete}: finished(file,complete) case 1n+rest Pending{file,storage,+offset,+remaining}: +count = Nat.min(U32.to_nat(chunk()),remaining) do IO<Bool>: next : Writing <- write_chunk(file,read_chunk(count,offset,storage,[]),count,remaining,offset) write_all(rest,next)def unsaved(pair: Tensor() & List<&2,U32>,operation: String,path: String,detail: String) -> Tensor(): (tensor,shape) = pair fault(operation,path,detail,shape)def written(complete: Bool,operation: String,path: String,tensor: Tensor()) -> Tensor(): match complete: case True{}: tensor case False{}: unsaved(Tensors.shape(~F32,tensor),operation,path,"cannot be written")# The prefix is written first; a raw file has none.def created(result: Result<&1,&1,U32 & String,File>,operation: String,+path: String,tensor: Tensor(),prefix: List<&2,U32>,storage: Array<F32>,+length: Nat) -> IO(Tensor()): match result: case Fail{error}: IO.pure(Tensor(),unsaved(Tensors.shape(~F32,tensor),operation,path,"cannot be created")) case Done{file}: do IO<Tensor()>: head : Written() <- File.write_bytes(file,prefix) complete : Bool <- write_all(Nat.add(length,2n),after_write(head,storage,length,0)) IO.pure(Tensor(),written(complete,operation,path,tensor))def saving(operation: String,+path: String,tensor: Tensor(),prefix: List<&2,U32>,buffer: Result<&2,&1,Tensors.Problem,Buffers.Buffer<F32>>) -> IO(Tensor()): match buffer: case Fail{problem}: IO.pure(Tensor(),tensor) case Done{Buffers.Buffer{storage,length,certificate}}: do IO<Tensor()>: result : Result<&1,&1,U32 & String,File> <- File.open(path,"w") created(result,operation,path,tensor,prefix,storage,length)# Writes the values in row-major order and returns the tensor. A failed tensor# writes nothing and keeps its problem.def save(+path: String,tensor: Tensor()) -> IO(Tensor()): Tensors.then(~Tensor(),~Tensor(),~IO(Tensor()),Layout.copy(~F32,tensor),kept => copy => saving("save",path,kept,[],Tensors.into_storage(~F32,copy)))def joined(shape: List<&2,U32>) -> String: match shape: case []: "" case [extent]: U32.show(extent) case extent <> rest: U32.show(extent) ++ ", " ++ joined(rest)# NumPy writes (), (3,) and (2, 3).def tuple(shape: List<&2,U32>) -> String: match shape: case []: "()" case [extent]: "(" ++ U32.show(extent) ++ ",)" case extent <> rest: "(" ++ joined(extent <> rest) ++ ")"def spaces(count: Nat) -> List<&2,U32>: match count: case 0n: [10] case 1n+rest: 32 <> spaces(rest)# Version 1: the header ends in a newline and the data starts at a multiple of 64 bytes.def framed(+text: List<&2,U32>,+length: U32) -> List<&2,U32>: +padding = ((64 - (11 + length) % 64) % 64 : U32) +total = (length + padding + 1 : U32) List.append(&2,U32,magic(),1 <> 0 <> (total % 256 : U32) <> (total / 256 : U32) <> List.append(&2,U32,text,spaces(U32.to_nat(padding))))def npy_header(shape: List<&2,U32>) -> List<&2,U32>: +text = codes("{'descr': '<f4', 'fortran_order': False, 'shape': " ++ tuple(shape) ++ ", }") framed(text,U32.from_nat(List.length(&2,U32,text)))def npy_saving(+path: String,pair: Tensor() & Tensor(),+shape: List<&2,U32>) -> IO(Tensor()): (kept,copy) = pair saving("save_npy",path,kept,npy_header(shape),Tensors.into_storage(~F32,copy))# Writes a .npy file that NumPy loads with the tensor's shape.def save_npy(+path: String,tensor: Tensor()) -> IO(Tensor()): Tensors.then(~Tensor(),~List<&2,U32>,~IO(Tensor()),Tensors.shape(~F32,tensor),tensor => shape => npy_saving(path,Layout.copy(~F32,tensor),shape))