def test_tmap()->None:
x:TMap[int,str]=TMap({3:'3',4:'4'})
def _f(a:str)->str:
return a
assert _f(x[3])=='3'
cache:Dict[TMap[int,str],bool]={}
cache[x] = True # Check hashability
assert cache[x] is True
def _make(b:Build):
build_setoutpaths(b, 1)
WLIB = mkwlib(lib_impl, Wdef)
IMEMs = [json2imem(j) for j in readjson(mklens(b).inputs.syspath)]
print(f"Inputs: {IMEMs}")
i = 0
# acc:List[Expr] = []
g = genexpr(WLIB, IMEMs)
written_bytes = 0
written_items = 0
time_start = time()
acci=set()
hb=time()
with open(mklens(b).out_examples.syspath,'wb') as f:
_add=examples2fd(f)
while time()<time_start+mklens(b).time2run_sec.val and \
written_items<mklens(b).maxitems.val:
# gt0=time()
r,mem,imems,exprw = next(g)
# gt1=time()
# print('gen time', gt1-gt0)
assert isinstance(imems[0][r], IVal), f"{imems[0][r]}"
i += 1
gi = 0
gexpr = gengather(r,mem)
# gg0=time()
exprs=[]
for expr in gexpr:
exprs.append(expr)
if gather_depth is not None:
exprs=exprs[-gather_depth:]
else:
exprs=[exprs[randint(0,len(exprs)-1)]]
for expr in exprs:
er=extrefs(expr)
ds=decls(expr)
# acc.append(expr)
for j in range(len(IMEMs)):
if len(ds)>0:
# print(gi, j, list(inps.keys()), print_expr(expr))
inps:IMem = TMap({k:imems[j][k] for k in er})
acci |= set(inps.values())
written_bytes+=_add(Example(inps,expr,imems[j][r]))
written_items+=1
hb2=time()
if written_items%100 == 0:
print(f".. NW {written_items} W {exprw[r]} DEP {depth(expr)} "
f"LAST_REF {r} WRBYTES {written_bytes // (1024) }K "
f"INPSZ {len(acci)} TIME {hb2-hb} "
f"VM {virtual_memory().used // 1024 // 1024}M")
hb=hb2
gi+=1
# Ident('concat'): lam('concat_a', lambda a:
# lam('concat_b', lambda b:
# intrin('concat', {'a':a,'b':b}))),
# # https://numpy.org/doc/stable/reference/generated/numpy.split.html
# Ident('split'): lam('split_a', lambda a:
# intrin('split', {'a':a})),
# }
def _transpose(args):
pass
def _concat(args):
pass
def _split(args):
pass
def mn(n: str) -> MethodName:
return MethodName(n)
lib: Lib = TMap({
mn('transpose'): LibEntry(mn('transpose'), ['a'], _transpose),
mn('concat'): LibEntry(mn('concat'), ['a', 'b'], _concat),
mn('split'): LibEntry(mn('split'), ['a'], _split)
})
def intrin(name:MethodName, args:List[Tuple[str,Expr]])->Expr:
acc = OrderedDict()
for (k,v) in args:
acc[str(k)] = v
return Intrin(name, TMap(acc))
def _succ(a) -> IExpr:
return IVal(a + 1)
def _prec(a) -> IExpr:
return IVal(a - 1)
def unpack(f):
def _f(args):
return f(**{an: a.val for an, a in args.items()})
return _f
def mn(n: str) -> MethodName:
return MethodName(n)
lib: Lib = TMap({
mn('add'): LibEntry(mn('add'), ['a', 'b'], unpack(_add)),
mn('mul'): LibEntry(mn('mul'), ['a', 'b'], unpack(_mul)),
mn('div'): LibEntry(mn('div'), ['a', 'b'], unpack(_div)),
mn('neg'): LibEntry(mn('neg'), ['a'], unpack(_neg)),
mn('sqr'): LibEntry(mn('sqr'), ['a'], unpack(_sqr)),
mn('sqrt'): LibEntry(mn('sqrt'), ['a'], unpack(_sqrt)),
mn('succ'): LibEntry(mn('succ'), ['a'], unpack(_succ)),
mn('prec'): LibEntry(mn('prec'), ['a'], unpack(_prec)),
})
def json2imem(d: dict) -> IMem:
return TMap({Ref(k): json2iexpr(v) for k, v in d.items()})
def _test_tmap_pickle(i:int)->TMap[int,str]:
return TMap({i:str(i)})
def bin2imem(d: BIN) -> IMem:
assert d.tag == Tag.imem, f"Unexpected tag {d.tag}"
return TMap({
Ref(str(i.tuple.v1.string)): bin2iexpr(i.tuple.v2.node)
for i in d.value.list.list
})
def genexpr(
wlib: WLib, inputs: List[IMem]
) -> Iterator[Tuple[Ref, TMap[Ref, Expr], List[IMem], Dict[Ref, int]]]:
""" Iterate over space of lambda-expressions with `len(inputs[0])` input
arguments. For every expression visited, provide results of
it's evaluation on every input of the `intputs` list.
Arguments:
* `wlib`: Weighted library of primitive operations. See also `mkwlib`.
* `inputs`: Collection of the inputs on which to grow the expression. All
inputs in list should be of the same size and use same names.
Yields a tuple of:
* Top-level reference `Ref`
* Map of `Ref -> Intrin`. Where `Intrin`s may contain more refs from the same
map.
* List of output expressions. Size of this list is equal to the size of
list of `inputs`.
* Current weight of the expression.
"""
# All inputs should provide the same input names
assert all([i.keys() == inputs[0].keys() for i in inputs])
nbatch = len(inputs)
lib = {k: wl[0] for k, wl in wlib.items()}
libws = {k: wl[1] for k, wl in wlib.items()}
exprcache: Dict[Ref, Expr] = {}
exprw: Dict[Ref, int] = {k: 1 for k in inputs[0].keys()}
valcache: List[Dict[Ref, IExpr]] = [OrderedDict(i.dict) for i in inputs]
W = 0
while True:
W += 1
for op in lib.values():
w = libws[op.name]
nargs = len(op.argnames)
vws: List[Tuple[Ref, int]] = list(exprw.items())
for valindices in permute(weights=[a[1] for a in vws],
nargs=nargs,
target_weight=W - w):
argrefs: List[Ref] = [vws[i][0] for i in valindices]
assert len(op.argnames) == len(argrefs)
e2name = Ref(mkname('val'))
e2expr = intrin(op.name,
[(nm, Val(ai))
for nm, ai in zip(op.argnames, argrefs)])
# TODO: Make this block customizable via callbacks
err = False
acc: List[IExpr] = []
for b in range(nbatch):
e2val, _ = interp(e2expr, TMap(lib), TMap(valcache[b]))
if isinstance(e2val, IError):
err = True
break
if isinstance(e2val, IVal) and isinstance(e2val.val, int):
if abs(e2val.val) > 10000 or abs(e2val.val) < -10000:
err = True
break
acc.append(e2val)
if err:
continue
for b in range(nbatch):
valcache[b][e2name] = acc[b]
exprcache[e2name] = e2expr
exprw[e2name] = W
yield (e2name, TMap(exprcache), [TMap(fd)
for fd in valcache], exprw)