add dreamcoder codebase

dreamcoder/__init__.py

@@ -0,0 +1,107 @@
+"""
+EC codebase Python library (AKA the "frontend")
+
+Module mapping details:
+
+TODO: remove module mapping code when backwards-compatibility is no longer required.
+
+The below module mapping is required for backwards-compatibility with old pickle files
+generated from before the EC codebase refactor. New files added to the codebase do not
+need to be added to the mapping, but if the existing modules are moved, then this the
+mapping needs to be updated to reflect the move or rename.
+
+The mapping uses the following pattern:
+
+    sys.modules[<old module path>] = <new module reference>
+
+This is because the previous structure of the codebase was completely flat, and when refactoring
+to a hierarchical files, loading previous pickle files no longer works properly. It is important
+to retain the ability to read old pickle files generated from official experiments. As a workaround,
+the old module paths are included below. A preferable alternative would be to export program state
+into JSON files instead of pickle files to avoid issues where the underlying classes change, so that
+could be a future improvement to this project. Until then, we use the module mapping workaround.
+
+For more info, see this StackOverflow answer: https://stackoverflow.com/a/2121918/2573242
+"""
+import sys
+
+from dreamcoder import differentiation
+from dreamcoder import dreamcoder
+from dreamcoder import enumeration
+from dreamcoder import fragmentGrammar
+from dreamcoder import fragmentUtilities
+from dreamcoder import frontier
+from dreamcoder import grammar
+from dreamcoder import likelihoodModel
+from dreamcoder import program
+from dreamcoder import primitiveGraph
+try:
+    from dreamcoder import recognition
+except:
+    print("Failure loading recognition - only acceptable if using pypy ",file=sys.stderr)
+from dreamcoder import task
+from dreamcoder import taskBatcher
+from dreamcoder import type
+from dreamcoder import utilities
+from dreamcoder import vs
+from dreamcoder.domains.misc import algolispPrimitives, deepcoderPrimitives
+from dreamcoder.domains.misc import RobustFillPrimitives
+from dreamcoder.domains.misc import napsPrimitives
+from dreamcoder.domains.tower import makeTowerTasks
+from dreamcoder.domains.tower import towerPrimitives
+from dreamcoder.domains.tower import tower_common
+from dreamcoder.domains.tower import main as tower_main
+from dreamcoder.domains.regex import groundtruthRegexes
+from dreamcoder.domains.regex import regexPrimitives
+from dreamcoder.domains.regex import makeRegexTasks
+#from dreamcoder.domains.regex import main as regex_main
+from dreamcoder.domains.logo import logoPrimitives
+from dreamcoder.domains.logo import makeLogoTasks
+from dreamcoder.domains.logo import main as logo_main
+from dreamcoder.domains.list import listPrimitives
+from dreamcoder.domains.list import makeListTasks
+from dreamcoder.domains.list import main as list_main
+from dreamcoder.domains.arithmetic import arithmeticPrimitives
+from dreamcoder.domains.text import textPrimitives
+from dreamcoder.domains.text import makeTextTasks
+from dreamcoder.domains.text import main as text_main
+
+sys.modules['differentiation'] = differentiation
+sys.modules['ec'] = dreamcoder
+sys.modules['enumeration'] = enumeration
+sys.modules['fragmentGrammar'] = fragmentGrammar
+sys.modules['fragmentUtilities'] = fragmentUtilities
+sys.modules['frontier'] = frontier
+sys.modules['grammar'] = grammar
+sys.modules['likelihoodModel'] = likelihoodModel
+sys.modules['program'] = program
+try: sys.modules['recognition'] = recognition
+except: pass
+sys.modules['task'] = task
+sys.modules['taskBatcher'] = taskBatcher
+sys.modules['type'] = type
+sys.modules['utilities'] = utilities
+sys.modules['vs'] = vs
+sys.modules['algolispPrimitives'] = algolispPrimitives
+sys.modules['RobustFillPrimitives'] = RobustFillPrimitives
+sys.modules['napsPrimitives'] = napsPrimitives
+sys.modules['makeTowerTasks'] = makeTowerTasks
+sys.modules['towerPrimitives'] = towerPrimitives
+sys.modules['tower_common'] = tower_common
+#sys.modules['tower'] = tower_main
+sys.modules['groundtruthRegexes'] = groundtruthRegexes
+sys.modules['regexPrimitives'] = regexPrimitives
+sys.modules['makeRegexTasks'] = makeRegexTasks
+#sys.modules['regexes'] = regex_main
+sys.modules['deepcoderPrimitives'] = deepcoderPrimitives
+sys.modules['logoPrimitives'] = logoPrimitives
+sys.modules['makeLogoTasks'] = makeLogoTasks
+#sys.modules['logo'] = logo_main
+sys.modules['listPrimitives'] = listPrimitives
+sys.modules['makeListTasks'] = makeListTasks
+#sys.modules['list'] = list_main
+sys.modules['arithmeticPrimitives'] = arithmeticPrimitives
+sys.modules['textPrimitives'] = textPrimitives
+sys.modules['makeTextTasks'] = makeTextTasks
+#sys.modules['text'] = text_main
+sys.modules['primitiveGraph'] = primitiveGraph

dreamcoder/compression.py

@@ -0,0 +1,282 @@
+import datetime
+import json
+import os
+import pickle
+import subprocess
+import sys
+
+from dreamcoder.fragmentGrammar import FragmentGrammar
+from dreamcoder.frontier import Frontier, FrontierEntry
+from dreamcoder.grammar import Grammar
+from dreamcoder.task import Task
+from dreamcoder.program import Program, Invented
+from dreamcoder.utilities import eprint, timing, callCompiled, get_root_dir
+from dreamcoder.vs import induceGrammar_Beta
+
+
+def induceGrammar(*args, **kwargs):
+    if sum(not f.empty for f in args[1]) == 0:
+        eprint("No nonempty frontiers, exiting grammar induction early.")
+        return args[0], args[1]
+    backend = kwargs.pop("backend", "pypy")
+    if 'pypy' in backend:
+        # pypy might not like some of the imports needed for the primitives
+        # but the primitive values are irrelevant for compression
+        # therefore strip them out and then replace them once we are done
+        # ditto for task data
+        g0,frontiers = args[0].strip_primitive_values(), \
+                       [front.strip_primitive_values() for front in args[1]]
+        original_tasks = {f.task.name: f.task for f in frontiers}
+        frontiers = [Frontier(f.entries, Task(f.task.name,f.task.request,[]))
+                     for f in frontiers ]
+        args = [g0,frontiers]
+
+    
+    with timing("Induced a grammar"):
+        if backend == "pypy":
+            g, newFrontiers = callCompiled(pypyInduce, *args, **kwargs)
+        elif backend == "rust":
+            g, newFrontiers = rustInduce(*args, **kwargs)
+        elif backend == "vs":
+            g, newFrontiers = rustInduce(*args, vs=True, **kwargs)
+        elif backend == "pypy_vs":
+            kwargs.pop('iteration')
+            kwargs.pop('topk_use_only_likelihood')
+            fn = '/tmp/vs.pickle'
+            with open(fn, 'wb') as handle:
+                pickle.dump((args, kwargs), handle)
+            eprint("For debugging purposes, the version space compression invocation has been saved to", fn)
+            g, newFrontiers = callCompiled(induceGrammar_Beta, *args, **kwargs)
+        elif backend == "ocaml":
+            kwargs.pop('iteration')
+            kwargs.pop('topk_use_only_likelihood')
+            kwargs['topI'] = 300
+            kwargs['bs'] = 1000000
+            g, newFrontiers = ocamlInduce(*args, **kwargs)
+        elif backend == "memorize":
+            g, newFrontiers = memorizeInduce(*args, **kwargs)
+        else:
+            assert False, "unknown compressor"
+
+    if 'pypy' in backend:
+        g, newFrontiers = g.unstrip_primitive_values(), \
+                          [front.unstrip_primitive_values() for front in newFrontiers]
+        newFrontiers = [Frontier(f.entries, original_tasks[f.task.name])
+                        for f in newFrontiers] 
+        
+
+    return g, newFrontiers
+
+def memorizeInduce(g, frontiers, **kwargs):
+    existingInventions = {p.uncurry()
+                          for p in g.primitives }
+    programs = {f.bestPosterior.program for f in frontiers if not f.empty}
+    newInventions = programs - existingInventions
+    newGrammar = Grammar.uniform([p for p in g.primitives] + \
+                                 [Invented(ni) for ni in newInventions])
+    
+    # rewrite in terms of new primitives
+    def substitute(p):
+        nonlocal newInventions
+        if p in newInventions: return Invented(p).uncurry()
+        return p
+    newFrontiers = [Frontier([FrontierEntry(program=np,
+                                            logPrior=newGrammar.logLikelihood(f.task.request, np),
+                                            logLikelihood=e.logLikelihood)
+                           for e in f
+                           for np in [substitute(e.program)] ],
+                             task=f.task)
+                 for f in frontiers ]
+    return newGrammar, newFrontiers
+    
+    
+        
+    
+
+def pypyInduce(*args, **kwargs):
+    kwargs.pop('iteration')
+    return FragmentGrammar.induceFromFrontiers(*args, **kwargs)
+
+
+def ocamlInduce(g, frontiers, _=None,
+                topK=1, pseudoCounts=1.0, aic=1.0,
+                structurePenalty=0.001, a=0, CPUs=1,
+                bs=1000000, topI=300):
+    # This is a dirty hack!
+    # Memory consumption increases with the number of CPUs
+    # And early on we have a lot of stuff to compress
+    # If this is the first iteration, only use a fraction of the available CPUs
+    if all(not p.isInvented for p in g.primitives):
+        if a > 3:
+            CPUs = max(1, int(CPUs / 6))
+        else:
+            CPUs = max(1, int(CPUs / 3))
+    else:
+        CPUs = max(1, int(CPUs / 2))
+    CPUs = 2
+
+    # X X X FIXME X X X
+    # for unknown reasons doing compression all in one go works correctly and doing it with Python and the outer loop causes problems
+    iterations = 99  # maximum number of components to add at once
+
+    while True:
+        g0 = g
+
+        originalFrontiers = frontiers
+        t2f = {f.task: f for f in frontiers}
+        frontiers = [f for f in frontiers if not f.empty]
+        message = {"arity": a,
+                   "topK": topK,
+                   "pseudoCounts": float(pseudoCounts),
+                   "aic": aic,
+                   "bs": bs,
+                   "topI": topI,
+                   "structurePenalty": float(structurePenalty),
+                   "CPUs": CPUs,
+                   "DSL": g.json(),
+                   "iterations": iterations,
+                   "frontiers": [f.json()
+                                 for f in frontiers]}
+
+        message = json.dumps(message)
+        if True:
+            timestamp = datetime.datetime.now().isoformat()
+            os.system("mkdir  -p compressionMessages")
+            fn = "compressionMessages/%s" % timestamp
+            with open(fn, "w") as f:
+                f.write(message)
+            eprint("Compression message saved to:", fn)
+
+        try:
+            # Get relative path
+            compressor_file = os.path.join(get_root_dir(), 'compression')
+            process = subprocess.Popen(compressor_file,
+                                       stdin=subprocess.PIPE,
+                                       stdout=subprocess.PIPE)
+            response, error = process.communicate(bytes(message, encoding="utf-8"))
+            response = json.loads(response.decode("utf-8"))
+        except OSError as exc:
+            raise exc
+
+        g = response["DSL"]
+        g = Grammar(g["logVariable"],
+                    [(l, p.infer(), p)
+                     for production in g["productions"]
+                     for l in [production["logProbability"]]
+                     for p in [Program.parse(production["expression"])]],
+                    continuationType=g0.continuationType)
+
+        frontiers = {original.task:
+                         Frontier([FrontierEntry(p,
+                                                 logLikelihood=e["logLikelihood"],
+                                                 logPrior=g.logLikelihood(original.task.request, p))
+                                   for e in new["programs"]
+                                   for p in [Program.parse(e["program"])]],
+                                  task=original.task)
+                     for original, new in zip(frontiers, response["frontiers"])}
+        frontiers = [frontiers.get(f.task, t2f[f.task])
+                     for f in originalFrontiers]
+        if iterations == 1 and len(g) > len(g0):
+            eprint("Grammar changed - running another round of consolidation.")
+            continue
+        else:
+            eprint("Finished consolidation.")
+            return g, frontiers
+
+
+def rustInduce(g0, frontiers, _=None,
+               topK=1, pseudoCounts=1.0, aic=1.0,
+               structurePenalty=0.001, a=0, CPUs=1, iteration=-1,
+               topk_use_only_likelihood=False,
+               vs=False):
+    def finite_logp(l):
+        return l if l != float("-inf") else -1000
+
+    message = {
+        "strategy": {"version-spaces": {"top_i": 50}}
+        if vs else
+        {"fragment-grammars": {}},
+        "params": {
+            "structure_penalty": structurePenalty,
+            "pseudocounts": int(pseudoCounts + 0.5),
+            "topk": topK,
+            "topk_use_only_likelihood": topk_use_only_likelihood,
+            "aic": aic if aic != float("inf") else None,
+            "arity": a,
+        },
+        "primitives": [{"name": p.name, "tp": str(t), "logp": finite_logp(l)}
+                       for l, t, p in g0.productions if p.isPrimitive],
+        "inventions": [{"expression": str(p.body),
+                        "logp": finite_logp(l)}  # -inf=-100
+                       for l, t, p in g0.productions if p.isInvented],
+        "variable_logprob": finite_logp(g0.logVariable),
+        "frontiers": [{
+            "task_tp": str(f.task.request),
+            "solutions": [{
+                "expression": str(e.program),
+                "logprior": finite_logp(e.logPrior),
+                "loglikelihood": e.logLikelihood,
+            } for e in f],
+        } for f in frontiers],
+    }
+
+    eprint("running rust compressor")
+
+    messageJson = json.dumps(message)
+
+    with open("jsonDebug", "w") as f:
+        f.write(messageJson)
+
+    # check which version of python we are using
+    # if >=3.6 do:
+    if sys.version_info[1] >= 6:
+        p = subprocess.Popen(
+            ['./rust_compressor/rust_compressor'],
+            encoding='utf-8',
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE)
+    elif sys.version_info[1] == 5:
+        p = subprocess.Popen(
+            ['./rust_compressor/rust_compressor'],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE)
+
+        messageJson = bytearray(messageJson, encoding='utf-8')
+        # convert messageJson string to bytes
+    else:
+        eprint("must be python 3.5 or 3.6")
+        assert False
+
+    p.stdin.write(messageJson)
+    p.stdin.flush()
+    p.stdin.close()
+
+    if p.returncode is not None:
+        raise ValueError("rust compressor failed")
+
+    if sys.version_info[1] >= 6:
+        resp = json.load(p.stdout)
+    elif sys.version_info[1] == 5:
+        import codecs
+        resp = json.load(codecs.getreader('utf-8')(p.stdout))
+
+    productions = [(x["logp"], p) for p, x in
+                   zip((p for (_, _, p) in g0.productions if p.isPrimitive), resp["primitives"])] + \
+                  [(i["logp"], Invented(Program.parse(i["expression"])))
+                   for i in resp["inventions"]]
+    productions = [(l if l is not None else float("-inf"), p)
+                   for l, p in productions]
+    g = Grammar.fromProductions(productions, resp["variable_logprob"], continuationType=g0.continuationType)
+    newFrontiers = [
+        Frontier(
+            [
+                FrontierEntry(
+                    Program.parse(
+                        s["expression"]),
+                    logPrior=s["logprior"],
+                    logLikelihood=s["loglikelihood"]) for s in r["solutions"]],
+            f.task) for f,
+                        r in zip(
+            frontiers,
+            resp["frontiers"])]
+    return g, newFrontiers

dreamcoder/deprecated/network.py

@@ -0,0 +1,479 @@
+"""
+Deprecated network.py module. This file only exists to support backwards-compatibility
+with old pickle files. See lib/__init__.py for more information.
+"""
+
+from __future__ import print_function
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from torch.nn.parameter import Parameter
+
+
+# UPGRADING TO INPUT -> OUTPUT -> TARGET
+# Todo:
+# [X] Output attending to input
+# [X] Target attending to output
+# [ ] check passing hidden state between encoders/decoder (+ pass c?)
+# [ ] add v_output
+
+
+def choose(matrix, idxs):
+    if isinstance(idxs, Variable):
+        idxs = idxs.data
+    assert(matrix.ndimension() == 2)
+    unrolled_idxs = idxs + \
+        torch.arange(0, matrix.size(0)).type_as(idxs) * matrix.size(1)
+    return matrix.view(matrix.nelement())[unrolled_idxs]
+
+
+class Network(nn.Module):
+    """
+    Todo:
+    - Beam search
+    - check if this is right? attend during P->FC rather than during softmax->P?
+    - allow length 0 inputs/targets
+    - give n_examples as input to FC
+    - Initialise new weights randomly, rather than as zeroes
+    """
+
+    def __init__(
+            self,
+            input_vocabulary,
+            target_vocabulary,
+            hidden_size=512,
+            embedding_size=128,
+            cell_type="LSTM"):
+        """
+        :param list input_vocabulary: list of possible inputs
+        :param list target_vocabulary: list of possible targets
+        """
+        super(Network, self).__init__()
+        self.h_input_encoder_size = hidden_size
+        self.h_output_encoder_size = hidden_size
+        self.h_decoder_size = hidden_size
+        self.embedding_size = embedding_size
+        self.input_vocabulary = input_vocabulary
+        self.target_vocabulary = target_vocabulary
+        # Number of tokens in input vocabulary
+        self.v_input = len(input_vocabulary)
+        # Number of tokens in target vocabulary
+        self.v_target = len(target_vocabulary)
+
+        self.cell_type = cell_type
+        if cell_type == 'GRU':
+            self.input_encoder_cell = nn.GRUCell(
+                input_size=self.v_input + 1,
+                hidden_size=self.h_input_encoder_size,
+                bias=True)
+            self.input_encoder_init = Parameter(
+                torch.rand(1, self.h_input_encoder_size))
+            self.output_encoder_cell = nn.GRUCell(
+                input_size=self.v_input +
+                1 +
+                self.h_input_encoder_size,
+                hidden_size=self.h_output_encoder_size,
+                bias=True)
+            self.decoder_cell = nn.GRUCell(
+                input_size=self.v_target + 1,
+                hidden_size=self.h_decoder_size,
+                bias=True)
+        if cell_type == 'LSTM':
+            self.input_encoder_cell = nn.LSTMCell(
+                input_size=self.v_input + 1,
+                hidden_size=self.h_input_encoder_size,
+                bias=True)
+            self.input_encoder_init = nn.ParameterList([Parameter(torch.rand(
+                1, self.h_input_encoder_size)), Parameter(torch.rand(1, self.h_input_encoder_size))])
+            self.output_encoder_cell = nn.LSTMCell(
+                input_size=self.v_input +
+                1 +
+                self.h_input_encoder_size,
+                hidden_size=self.h_output_encoder_size,
+                bias=True)
+            self.output_encoder_init_c = Parameter(
+                torch.rand(1, self.h_output_encoder_size))
+            self.decoder_cell = nn.LSTMCell(
+                input_size=self.v_target + 1,
+                hidden_size=self.h_decoder_size,
+                bias=True)
+            self.decoder_init_c = Parameter(torch.rand(1, self.h_decoder_size))
+
+        self.W = nn.Linear(
+            self.h_output_encoder_size +
+            self.h_decoder_size,
+            self.embedding_size)
+        self.V = nn.Linear(self.embedding_size, self.v_target + 1)
+        self.input_A = nn.Bilinear(
+            self.h_input_encoder_size,
+            self.h_output_encoder_size,
+            1,
+            bias=False)
+        self.output_A = nn.Bilinear(
+            self.h_output_encoder_size,
+            self.h_decoder_size,
+            1,
+            bias=False)
+        self.input_EOS = torch.zeros(1, self.v_input + 1)
+        self.input_EOS[:, -1] = 1
+        self.input_EOS = Parameter(self.input_EOS)
+        self.output_EOS = torch.zeros(1, self.v_input + 1)
+        self.output_EOS[:, -1] = 1
+        self.output_EOS = Parameter(self.output_EOS)
+        self.target_EOS = torch.zeros(1, self.v_target + 1)
+        self.target_EOS[:, -1] = 1
+        self.target_EOS = Parameter(self.target_EOS)
+
+    def __getstate__(self):
+        if hasattr(self, 'opt'):
+            return dict([(k, v) for k, v in self.__dict__.items(
+            ) if k is not 'opt'] + [('optstate', self.opt.state_dict())])
+            # return {**{k:v for k,v in self.__dict__.items() if k is not 'opt'},
+            #         'optstate': self.opt.state_dict()}
+        else:
+            return self.__dict__
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        # Legacy:
+        if isinstance(self.input_encoder_init, tuple):
+            self.input_encoder_init = nn.ParameterList(
+                list(self.input_encoder_init))
+
+    def clear_optimiser(self):
+        if hasattr(self, 'opt'):
+            del self.opt
+        if hasattr(self, 'optstate'):
+            del self.optstate
+
+    def get_optimiser(self):
+        self.opt = torch.optim.Adam(self.parameters(), lr=0.001)
+        if hasattr(self, 'optstate'):
+            self.opt.load_state_dict(self.optstate)
+
+    def optimiser_step(self, inputs, outputs, target):
+        if not hasattr(self, 'opt'):
+            self.get_optimiser()
+        score = self.score(inputs, outputs, target, autograd=True).mean()
+        (-score).backward()
+        self.opt.step()
+        self.opt.zero_grad()
+        return score.data[0]
+
+    def set_target_vocabulary(self, target_vocabulary):
+        if target_vocabulary == self.target_vocabulary:
+            return
+
+        V_weight = []
+        V_bias = []
+        decoder_ih = []
+
+        for i in range(len(target_vocabulary)):
+            if target_vocabulary[i] in self.target_vocabulary:
+                j = self.target_vocabulary.index(target_vocabulary[i])
+                V_weight.append(self.V.weight.data[j:j + 1])
+                V_bias.append(self.V.bias.data[j:j + 1])
+                decoder_ih.append(self.decoder_cell.weight_ih.data[:, j:j + 1])
+            else:
+                V_weight.append(torch.zeros(1, self.V.weight.size(1)))
+                V_bias.append(torch.ones(1) * -10)
+                decoder_ih.append(
+                    torch.zeros(
+                        self.decoder_cell.weight_ih.data.size(0), 1))
+
+        V_weight.append(self.V.weight.data[-1:])
+        V_bias.append(self.V.bias.data[-1:])
+        decoder_ih.append(self.decoder_cell.weight_ih.data[:, -1:])
+
+        self.target_vocabulary = target_vocabulary
+        self.v_target = len(target_vocabulary)
+        self.target_EOS.data = torch.zeros(1, self.v_target + 1)
+        self.target_EOS.data[:, -1] = 1
+
+        self.V.weight.data = torch.cat(V_weight, dim=0)
+        self.V.bias.data = torch.cat(V_bias, dim=0)
+        self.V.out_features = self.V.bias.data.size(0)
+
+        self.decoder_cell.weight_ih.data = torch.cat(decoder_ih, dim=1)
+        self.decoder_cell.input_size = self.decoder_cell.weight_ih.data.size(1)
+
+        self.clear_optimiser()
+
+    def input_encoder_get_init(self, batch_size):
+        if self.cell_type == "GRU":
+            return self.input_encoder_init.repeat(batch_size, 1)
+        if self.cell_type == "LSTM":
+            return tuple(x.repeat(batch_size, 1)
+                         for x in self.input_encoder_init)
+
+    def output_encoder_get_init(self, input_encoder_h):
+        if self.cell_type == "GRU":
+            return input_encoder_h
+        if self.cell_type == "LSTM":
+            return (
+                input_encoder_h,
+                self.output_encoder_init_c.repeat(
+                    input_encoder_h.size(0),
+                    1))
+
+    def decoder_get_init(self, output_encoder_h):
+        if self.cell_type == "GRU":
+            return output_encoder_h
+        if self.cell_type == "LSTM":
+            return (
+                output_encoder_h,
+                self.decoder_init_c.repeat(
+                    output_encoder_h.size(0),
+                    1))
+
+    def cell_get_h(self, cell_state):
+        if self.cell_type == "GRU":
+            return cell_state
+        if self.cell_type == "LSTM":
+            return cell_state[0]
+
+    def score(self, inputs, outputs, target, autograd=False):
+        inputs = self.inputsToTensors(inputs)
+        outputs = self.inputsToTensors(outputs)
+        target = self.targetToTensor(target)
+        target, score = self.run(inputs, outputs, target=target, mode="score")
+        # target = self.tensorToOutput(target)
+        if autograd:
+            return score
+        else:
+            return score.data
+
+    def sample(self, inputs, outputs):
+        inputs = self.inputsToTensors(inputs)
+        outputs = self.inputsToTensors(outputs)
+        target, score = self.run(inputs, outputs, mode="sample")
+        target = self.tensorToOutput(target)
+        return target
+
+    def sampleAndScore(self, inputs, outputs, nRepeats=None):
+        inputs = self.inputsToTensors(inputs)
+        outputs = self.inputsToTensors(outputs)
+        if nRepeats is None:
+            target, score = self.run(inputs, outputs, mode="sample")
+            target = self.tensorToOutput(target)
+            return target, score.data
+        else:
+            target = []
+            score = []
+            for i in range(nRepeats):
+                # print("repeat %d" % i)
+                t, s = self.run(inputs, outputs, mode="sample")
+                t = self.tensorToOutput(t)
+                target.extend(t)
+                score.extend(list(s.data))
+            return target, score
+
+    def run(self, inputs, outputs, target=None, mode="sample"):
+        """
+        :param mode: "score" returns log p(target|input), "sample" returns target ~ p(-|input)
+        :param List[LongTensor] inputs: n_examples * (max_length_input * batch_size)
+        :param List[LongTensor] target: max_length_target * batch_size
+        """
+        assert((mode == "score" and target is not None) or mode == "sample")
+
+        n_examples = len(inputs)
+        max_length_input = [inputs[j].size(0) for j in range(n_examples)]
+        max_length_output = [outputs[j].size(0) for j in range(n_examples)]
+        max_length_target = target.size(0) if target is not None else 10
+        batch_size = inputs[0].size(1)
+
+        score = Variable(torch.zeros(batch_size))
+        inputs_scatter = [Variable(torch.zeros(max_length_input[j], batch_size, self.v_input + 1).scatter_(
+            2, inputs[j][:, :, None], 1)) for j in range(n_examples)]  # n_examples * (max_length_input * batch_size * v_input+1)
+        outputs_scatter = [Variable(torch.zeros(max_length_output[j], batch_size, self.v_input + 1).scatter_(
+            2, outputs[j][:, :, None], 1)) for j in range(n_examples)]  # n_examples * (max_length_output * batch_size * v_input+1)
+        if target is not None:
+            target_scatter = Variable(torch.zeros(max_length_target,
+                                                  batch_size,
+                                                  self.v_target + 1).scatter_(2,
+                                                                              target[:,
+                                                                                     :,
+                                                                                     None],
+                                                                              1))  # max_length_target * batch_size * v_target+1
+
+        # -------------- Input Encoder -------------
+
+        # n_examples * (max_length_input * batch_size * h_encoder_size)
+        input_H = []
+        input_embeddings = []  # h for example at INPUT_EOS
+        # 0 until (and including) INPUT_EOS, then -inf
+        input_attention_mask = []
+        for j in range(n_examples):
+            active = torch.Tensor(max_length_input[j], batch_size).byte()
+            active[0, :] = 1
+            state = self.input_encoder_get_init(batch_size)
+            hs = []
+            for i in range(max_length_input[j]):
+                state = self.input_encoder_cell(
+                    inputs_scatter[j][i, :, :], state)
+                if i + 1 < max_length_input[j]:
+                    active[i + 1, :] = active[i, :] * \
+                        (inputs[j][i, :] != self.v_input)
+                h = self.cell_get_h(state)
+                hs.append(h[None, :, :])
+            input_H.append(torch.cat(hs, 0))
+            embedding_idx = active.sum(0).long() - 1
+            embedding = input_H[j].gather(0, Variable(
+                embedding_idx[None, :, None].repeat(1, 1, self.h_input_encoder_size)))[0]
+            input_embeddings.append(embedding)
+            input_attention_mask.append(Variable(active.float().log()))
+
+        # -------------- Output Encoder -------------
+
+        def input_attend(j, h_out):
+            """
+            'general' attention from https://arxiv.org/pdf/1508.04025.pdf
+            :param j: Index of example
+            :param h_out: batch_size * h_output_encoder_size
+            """
+            scores = self.input_A(
+                input_H[j].view(
+                    max_length_input[j] * batch_size,
+                    self.h_input_encoder_size),
+                h_out.view(
+                    batch_size,
+                    self.h_output_encoder_size).repeat(
+                    max_length_input[j],
+                    1)).view(
+                max_length_input[j],
+                batch_size) + input_attention_mask[j]
+            c = (F.softmax(scores[:, :, None], dim=0) * input_H[j]).sum(0)
+            return c
+
+        # n_examples * (max_length_input * batch_size * h_encoder_size)
+        output_H = []
+        output_embeddings = []  # h for example at INPUT_EOS
+        # 0 until (and including) INPUT_EOS, then -inf
+        output_attention_mask = []
+        for j in range(n_examples):
+            active = torch.Tensor(max_length_output[j], batch_size).byte()
+            active[0, :] = 1
+            state = self.output_encoder_get_init(input_embeddings[j])
+            hs = []
+            h = self.cell_get_h(state)
+            for i in range(max_length_output[j]):
+                state = self.output_encoder_cell(torch.cat(
+                    [outputs_scatter[j][i, :, :], input_attend(j, h)], 1), state)
+                if i + 1 < max_length_output[j]:
+                    active[i + 1, :] = active[i, :] * \
+                        (outputs[j][i, :] != self.v_input)
+                h = self.cell_get_h(state)
+                hs.append(h[None, :, :])
+            output_H.append(torch.cat(hs, 0))
+            embedding_idx = active.sum(0).long() - 1
+            embedding = output_H[j].gather(0, Variable(
+                embedding_idx[None, :, None].repeat(1, 1, self.h_output_encoder_size)))[0]
+            output_embeddings.append(embedding)
+            output_attention_mask.append(Variable(active.float().log()))
+
+        # ------------------ Decoder -----------------
+
+        def output_attend(j, h_dec):
+            """
+            'general' attention from https://arxiv.org/pdf/1508.04025.pdf
+            :param j: Index of example
+            :param h_dec: batch_size * h_decoder_size
+            """
+            scores = self.output_A(
+                output_H[j].view(
+                    max_length_output[j] * batch_size,
+                    self.h_output_encoder_size),
+                h_dec.view(
+                    batch_size,
+                    self.h_decoder_size).repeat(
+                    max_length_output[j],
+                    1)).view(
+                max_length_output[j],
+                batch_size) + output_attention_mask[j]
+            c = (F.softmax(scores[:, :, None], dim=0) * output_H[j]).sum(0)
+            return c
+
+        # Multi-example pooling: Figure 3, https://arxiv.org/pdf/1703.07469.pdf
+        target = target if mode == "score" else torch.zeros(
+            max_length_target, batch_size).long()
+        decoder_states = [
+            self.decoder_get_init(
+                output_embeddings[j]) for j in range(n_examples)]  # P
+        active = torch.ones(batch_size).byte()
+        for i in range(max_length_target):
+            FC = []
+            for j in range(n_examples):
+                h = self.cell_get_h(decoder_states[j])
+                p_aug = torch.cat([h, output_attend(j, h)], 1)
+                FC.append(F.tanh(self.W(p_aug)[None, :, :]))
+            # batch_size * embedding_size
+            m = torch.max(torch.cat(FC, 0), 0)[0]
+            logsoftmax = F.log_softmax(self.V(m), dim=1)
+            if mode == "sample":
+                target[i, :] = torch.multinomial(
+                    logsoftmax.data.exp(), 1)[:, 0]
+            score = score + \
+                choose(logsoftmax, target[i, :]) * Variable(active.float())
+            active *= (target[i, :] != self.v_target)
+            for j in range(n_examples):
+                if mode == "score":
+                    target_char_scatter = target_scatter[i, :, :]
+                elif mode == "sample":
+                    target_char_scatter = Variable(torch.zeros(
+                        batch_size, self.v_target + 1).scatter_(1, target[i, :, None], 1))
+                decoder_states[j] = self.decoder_cell(
+                    target_char_scatter, decoder_states[j])
+        return target, score
+
+    def inputsToTensors(self, inputss):
+        """
+        :param inputss: size = nBatch * nExamples
+        """
+        tensors = []
+        for j in range(len(inputss[0])):
+            inputs = [x[j] for x in inputss]
+            maxlen = max(len(s) for s in inputs)
+            t = torch.ones(
+                1 if maxlen == 0 else maxlen + 1,
+                len(inputs)).long() * self.v_input
+            for i in range(len(inputs)):
+                s = inputs[i]
+                if len(s) > 0:
+                    t[:len(s), i] = torch.LongTensor(
+                        [self.input_vocabulary.index(x) for x in s])
+            tensors.append(t)
+        return tensors
+
+    def targetToTensor(self, targets):
+        """
+        :param targets:
+        """
+        maxlen = max(len(s) for s in targets)
+        t = torch.ones(
+            1 if maxlen == 0 else maxlen + 1,
+            len(targets)).long() * self.v_target
+        for i in range(len(targets)):
+            s = targets[i]
+            if len(s) > 0:
+                t[:len(s), i] = torch.LongTensor(
+                    [self.target_vocabulary.index(x) for x in s])
+        return t
+
+    def tensorToOutput(self, tensor):
+        """
+        :param tensor: max_length * batch_size
+        """
+        out = []
+        for i in range(tensor.size(1)):
+            l = tensor[:, i].tolist()
+            if l[0] == self.v_target:
+                out.append([])
+            elif self.v_target in l:
+                final = tensor[:, i].tolist().index(self.v_target)
+                out.append([self.target_vocabulary[x]
+                            for x in tensor[:final, i]])
+            else:
+                out.append([self.target_vocabulary[x] for x in tensor[:, i]])
+        return out

dreamcoder/differentiation.py

@@ -0,0 +1,393 @@
+import math
+import random
+from dreamcoder.utilities import *
+
+
+class InvalidLoss(Exception):
+    pass
+
+
+class DN(object):
+    '''differentiable node: parent object of every differentiable operation'''
+
+    def __init__(self, arguments):
+        self.gradient = None
+        if arguments != []:
+            self.data = None
+        self.arguments = arguments
+
+        # descendents: every variable that takes this variable as input
+        # descendents: [(DN,float)]
+        # the additional float parameter is d Descendent / d This
+        self.descendents = []
+
+        self.recalculate()
+
+    def __str__(self):
+        if self.arguments == []:
+            return self.name
+        return "(%s %s)" % (self.name, " ".join(str(x)
+                                                for x in self.arguments))
+
+    def __repr__(self):
+        return "DN(op = %s, data = %s, grad = %s, #descendents = %d, args = %s)" % (
+            self.name, self.data, self.gradient, len(self.descendents), self.arguments)
+
+    @property
+    def derivative(self): return self.differentiate()
+
+    def differentiate(self):
+        if self.gradient is None:
+            self.gradient = sum(partial * descendent.differentiate()
+                                for descendent, partial in self.descendents)
+        return self.gradient
+
+    def zeroEverything(self):
+        if self.gradient is None and self.descendents == [] and (
+                self.data is None or self.arguments == []):
+            return
+
+        self.gradient = None
+        self.descendents = []
+        if self.arguments != []:
+            self.data = None
+
+        for x in self.arguments:
+            x.zeroEverything()
+
+    def lightweightRecalculate(self):
+        return self.forward(*[a.lightweightRecalculate()
+                              for a in self.arguments])
+
+    def recalculate(self):
+        if self.data is None:
+            inputs = [a.recalculate() for a in self.arguments]
+            self.data = self.forward(*inputs)
+            # if invalid(self.data):
+            #     eprint("I am invalid",repr(self))
+            #     eprint("Here are my inputs",inputs)
+            #     self.zeroEverything()
+            #     eprint("Here I am after being zeroed",repr(self))
+            #     raise Exception('invalid loss')
+            #assert valid(self.data)
+            partials = self.backward(*inputs)
+            for d, a in zip(partials, self.arguments):
+                # if invalid(d):
+                #     eprint("I have an invalid derivative",self)
+                #     eprint("Inputs",inputs)
+                #     eprint("partials",partials)
+                #     raise Exception('invalid derivative')
+                a.descendents.append((self, d))
+        return self.data
+
+    def backPropagation(self):
+        self.gradient = 1.
+        self.recursivelyDifferentiate()
+
+    def recursivelyDifferentiate(self):
+        self.differentiate()
+        for x in self.arguments:
+            x.recursivelyDifferentiate()
+
+    def updateNetwork(self):
+        self.zeroEverything()
+        l = self.recalculate()
+        self.backPropagation()
+        return l
+
+    def log(self): return Logarithm(self)
+
+    def square(self): return Square(self)
+
+    def exp(self): return Exponentiation(self)
+
+    def clamp(self, l, u): return Clamp(self, l, u)
+
+    def __abs__(self): return AbsoluteValue(self)
+
+    def __add__(self, o): return Addition(self, Placeholder.maybe(o))
+
+    def __radd__(self, o): return Addition(self, Placeholder.maybe(o))
+
+    def __sub__(self, o): return Subtraction(self, Placeholder.maybe(o))
+
+    def __rsub__(self, o): return Subtraction(Placeholder.maybe(o), self)
+
+    def __mul__(self, o): return Multiplication(self, Placeholder.maybe(o))
+
+    def __rmul__(self, o): return Multiplication(self, Placeholder.maybe(o))
+
+    def __neg__(self): return Negation(self)
+
+    def __truediv__(self, o): return Division(self, Placeholder.maybe(o))
+
+    def __rtruediv__(self, o): return Division(Placeholder.maybe(o), self)
+
+    def numericallyVerifyGradients(self, parameters):
+        calculatedGradients = [p.derivative for p in parameters]
+        e = 0.00001
+        for j, p in enumerate(parameters):
+            p.data -= e
+            y1 = self.lightweightRecalculate()
+            p.data += 2 * e
+            y2 = self.lightweightRecalculate()
+            p.data -= e
+            d = (y2 - y1) / (2 * e)
+            if abs(calculatedGradients[j] - d) > 0.1:
+                eprint(
+                    "Bad gradient: expected %f, got %f" %
+                    (d, calculatedGradients[j]))
+
+    def gradientDescent(
+            self,
+            parameters,
+            _=None,
+            lr=0.001,
+            steps=10**3,
+            update=None):
+        for j in range(steps):
+            l = self.updateNetwork()
+            if update is not None and j % update == 0:
+                eprint("LOSS:", l)
+                for p in parameters:
+                    eprint(p.data, '\t', p.derivative)
+            if invalid(l):
+                raise InvalidLoss()
+
+            for p in parameters:
+                p.data -= lr * p.derivative
+        return self.data
+
+    def restartingOptimize(self, parameters, _=None, attempts=1,
+                           s=1., decay=0.5, grow=0.1,
+                           lr=0.1, steps=10**3, update=None):
+        ls = []
+        for _ in range(attempts):
+            for p in parameters:
+                p.data = random.random()*10 - 5
+            ls.append(
+                self.resilientBackPropagation(
+                    parameters, lr=lr, steps=steps,
+                    decay=decay, grow=grow))
+        return min(ls)
+
+    def resilientBackPropagation(
+            self,
+            parameters,
+            _=None,
+            decay=0.5,
+            grow=1.2,
+            lr=0.1,
+            steps=10**3,
+            update=None):
+        previousSign = [None] * len(parameters)
+        lr = [lr] * len(parameters)
+        for j in range(steps):
+            l = self.updateNetwork()
+
+            if update is not None and j % update == 0:
+                eprint("LOSS:", l)
+                eprint("\t".join(str(p.derivative) for p in parameters))
+            if invalid(l):
+                raise InvalidLoss()
+
+            newSigns = [p.derivative > 0 for p in parameters]
+            for i, p in enumerate(parameters):
+                if p.derivative > 0:
+                    p.data -= lr[i]
+                elif p.derivative < 0:
+                    p.data += lr[i]
+                if previousSign[i] is not None:
+                    if previousSign[i] == newSigns[i]:
+                        lr[i] *= grow
+                    else:
+                        lr[i] *= decay
+            previousSign = newSigns
+
+        return self.data
+
+
+class Placeholder(DN):
+    COUNTER = 0
+
+    def __init__(self, initialValue=0., name=None):
+        self.data = initialValue
+        super(Placeholder, self).__init__([])
+        if name is None:
+            name = "p_" + str(Placeholder.COUNTER)
+            Placeholder.COUNTER += 1
+        self.name = name
+
+    @staticmethod
+    def named(namePrefix, initialValue=0.):
+        p = Placeholder(initialValue, namePrefix + str(Placeholder.COUNTER))
+        Placeholder.COUNTER += 1
+        return p
+
+    def __str__(self):
+        return "Placeholder(%s = %s)" % (self.name, self.data)
+
+    @staticmethod
+    def maybe(x):
+        if isinstance(x, DN):
+            return x
+        return Placeholder(float(x))
+
+    def forward(self): return self.data
+
+    def backward(self): return []
+
+
+class Clamp(DN):
+    def __init__(self, x, l, u):
+        assert u > l
+        self.l = l
+        self.u = u
+        super(Clamp, self).__init__([x])
+        self.name = "clamp"
+
+    def forward(self, x):
+        if x > self.u:
+            return self.u
+        if x < self.l:
+            return self.l
+        return x
+
+    def backward(self, x):
+        if x > self.u or x < self.l:
+            return [0.]
+        else:
+            return [1.]
+
+
+class Addition(DN):
+    def __init__(self, x, y):
+        super(Addition, self).__init__([x, y])
+        self.name = '+'
+
+    def forward(self, x, y): return x + y
+
+    def backward(self, x, y): return [1., 1.]
+
+
+class Subtraction(DN):
+    def __init__(self, x, y):
+        super(Subtraction, self).__init__([x, y])
+        self.name = '-'
+
+    def forward(self, x, y): return x - y
+
+    def backward(self, x, y): return [1., -1.]
+
+
+class Negation(DN):
+    def __init__(self, x):
+        super(Negation, self).__init__([x])
+        self.name = '-'
+
+    def forward(self, x): return -x
+
+    def backward(self, x): return [-1.]
+
+
+class AbsoluteValue(DN):
+    def __init__(self, x):
+        super(AbsoluteValue, self).__init__([x])
+        self.name = 'abs'
+
+    def forward(self, x): return abs(x)
+
+    def backward(self, x):
+        if x > 0:
+            return [1.]
+        return [-1.]
+
+
+class Multiplication(DN):
+    def __init__(self, x, y):
+        super(Multiplication, self).__init__([x, y])
+        self.name = '*'
+
+    def forward(self, x, y): return x * y
+
+    def backward(self, x, y): return [y, x]
+
+
+class Division(DN):
+    def __init__(self, x, y):
+        super(Division, self).__init__([x, y])
+        self.name = '/'
+
+    def forward(self, x, y): return x / y
+
+    def backward(self, x, y): return [1.0 / y, -x / (y * y)]
+
+
+class Square(DN):
+    def __init__(self, x):
+        super(Square, self).__init__([x])
+        self.name = 'sq'
+
+    def forward(self, x): return x * x
+
+    def backward(self, x): return [2 * x]
+
+
+class Exponentiation(DN):
+    def __init__(self, x):
+        super(Exponentiation, self).__init__([x])
+        self.name = 'exp'
+
+    def forward(self, x): return math.exp(x)
+
+    def backward(self, x): return [math.exp(x)]
+
+
+class Logarithm(DN):
+    def __init__(self, x):
+        super(Logarithm, self).__init__([x])
+        self.name = 'log'
+
+    def forward(self, x): return math.log(x)
+
+    def backward(self, x): return [1. / x]
+
+
+class LSE(DN):
+    def __init__(self, xs):
+        super(LSE, self).__init__(xs)
+        self.name = 'LSE'
+
+    def forward(self, *xs):
+        m = max(xs)
+        return m + math.log(sum(math.exp(y - m) for y in xs))
+
+    def backward(self, *xs):
+        m = max(xs)
+        zm = sum(math.exp(x - m) for x in xs)
+        return [math.exp(x - m) / zm for x in xs]
+
+
+if __name__ == "__main__":
+    x = Placeholder(10., "x")
+    y = Placeholder(2., "y")
+    z = x - LSE([x, y])
+    z.updateNetwork()
+    eprint("dL/dx = %f\tdL/dy = %f" % (x.derivative, y.derivative))
+
+    x.data = 2.
+    y.data = 10.
+    z.updateNetwork()
+    eprint("dL/dx = %f\tdL/dy = %f" % (x.differentiate(), y.differentiate()))
+
+    x.data = 2.
+    y.data = 2.
+    z.updateNetwork()
+    eprint("z = ", z.data, z)
+    eprint("dL/dx = %f\tdL/dy = %f" % (x.differentiate(), y.differentiate()))
+
+    loss = -z
+    eprint(loss)
+
+    lr = 0.001
+    loss.gradientDescent([x, y], steps=10000, update=1000)

dreamcoder/domains/arithmetic/arithmeticPrimitives.py

@@ -0,0 +1,58 @@
+from dreamcoder.program import *
+from dreamcoder.type import *
+
+
+def _addition(x): return lambda y: x + y
+
+
+def _subtraction(x): return lambda y: x - y
+
+
+def _division(x): return lambda y: x / y
+
+
+subtraction = Primitive("-",
+                        arrow(tint, arrow(tint, tint)),
+                        _subtraction)
+real_subtraction = Primitive("-.",
+                             arrow(treal, treal, treal),
+                             _subtraction)
+addition = Primitive("+",
+                     arrow(tint, arrow(tint, tint)),
+                     Curried(_addition))
+real_addition = Primitive("+.",
+                          arrow(treal, treal, treal),
+                          _addition)
+
+
+def _multiplication(x): return lambda y: x * y
+
+
+multiplication = Primitive("*",
+                           arrow(tint, arrow(tint, tint)),
+                           _multiplication)
+real_multiplication = Primitive("*.",
+                                arrow(treal, treal, treal),
+                                _multiplication)
+real_division = Primitive("/.",
+                          arrow(treal, treal, treal),
+                          _division)
+
+
+def _power(a): return lambda b: a**b
+
+
+real_power = Primitive("power",
+                       arrow(treal, treal, treal),
+                       _power)
+
+k1 = Primitive("1", tint, 1)
+k_negative1 = Primitive("negative_1", tint, -1)
+k0 = Primitive("0", tint, 0)
+for n in range(2,10):
+    Primitive(str(n),tint,n)
+
+f1 = Primitive("1.", treal, 1.)
+f0 = Primitive("0.", treal, 0)
+real = Primitive("REAL", treal, None)
+fpi = Primitive("pi", treal, 3.14)

dreamcoder/domains/list/listPrimitives.py

@@ -0,0 +1,546 @@
+from dreamcoder.program import Primitive, Program
+from dreamcoder.grammar import Grammar
+from dreamcoder.type import tlist, tint, tbool, arrow, t0, t1, t2
+
+import math
+from functools import reduce
+
+
+def _flatten(l): return [x for xs in l for x in xs]
+
+def _range(n):
+    if n < 100: return list(range(n))
+    raise ValueError()
+def _if(c): return lambda t: lambda f: t if c else f
+
+
+def _and(x): return lambda y: x and y
+
+
+def _or(x): return lambda y: x or y
+
+
+def _addition(x): return lambda y: x + y
+
+
+def _subtraction(x): return lambda y: x - y
+
+
+def _multiplication(x): return lambda y: x * y
+
+
+def _negate(x): return -x
+
+
+def _reverse(x): return list(reversed(x))
+
+
+def _append(x): return lambda y: x + y
+
+
+def _cons(x): return lambda y: [x] + y
+
+
+def _car(x): return x[0]
+
+
+def _cdr(x): return x[1:]
+
+
+def _isEmpty(x): return x == []
+
+
+def _single(x): return [x]
+
+
+def _slice(x): return lambda y: lambda l: l[x:y]
+
+
+def _map(f): return lambda l: list(map(f, l))
+
+
+def _zip(a): return lambda b: lambda f: list(map(lambda x,y: f(x)(y), a, b))
+
+
+def _mapi(f): return lambda l: list(map(lambda i_x: f(i_x[0])(i_x[1]), enumerate(l)))
+
+
+def _reduce(f): return lambda x0: lambda l: reduce(lambda a, x: f(a)(x), l, x0)
+
+
+def _reducei(f): return lambda x0: lambda l: reduce(
+    lambda a, t: f(t[0])(a)(t[1]), enumerate(l), x0)
+
+
+def _fold(l): return lambda x0: lambda f: reduce(
+    lambda a, x: f(x)(a), l[::-1], x0)
+
+
+def _eq(x): return lambda y: x == y
+
+
+def _eq0(x): return x == 0
+
+
+def _a1(x): return x + 1
+
+
+def _d1(x): return x - 1
+
+
+def _mod(x): return lambda y: x % y
+
+
+def _not(x): return not x
+
+
+def _gt(x): return lambda y: x > y
+
+
+def _index(j): return lambda l: l[j]
+
+
+def _replace(f): return lambda lnew: lambda lin: _flatten(
+    lnew if f(i)(x) else [x] for i, x in enumerate(lin))
+
+
+def _isPrime(n):
+    return n in {
+        2,
+        3,
+        5,
+        7,
+        11,
+        13,
+        17,
+        19,
+        23,
+        29,
+        31,
+        37,
+        41,
+        43,
+        47,
+        53,
+        59,
+        61,
+        67,
+        71,
+        73,
+        79,
+        83,
+        89,
+        97,
+        101,
+        103,
+        107,
+        109,
+        113,
+        127,
+        131,
+        137,
+        139,
+        149,
+        151,
+        157,
+        163,
+        167,
+        173,
+        179,
+        181,
+        191,
+        193,
+        197,
+        199}
+
+
+def _isSquare(n):
+    return int(math.sqrt(n)) ** 2 == n
+
+
+def _appendmap(f): lambda xs: [y for x in xs for y in f(x)]
+
+
+def _filter(f): return lambda l: list(filter(f, l))
+
+
+def _any(f): return lambda l: any(f(x) for x in l)
+
+
+def _all(f): return lambda l: all(f(x) for x in l)
+
+
+def _find(x):
+    def _inner(l):
+        try:
+            return l.index(x)
+        except ValueError:
+            return -1
+    return _inner
+
+
+def _unfold(x): return lambda p: lambda h: lambda n: __unfold(p, f, n, x)
+
+
+def __unfold(p, f, n, x, recursion_limit=50):
+    if recursion_limit <= 0:
+        raise RecursionDepthExceeded()
+    if p(x):
+        return []
+    return [f(x)] + __unfold(p, f, n, n(x), recursion_limit - 1)
+
+
+class RecursionDepthExceeded(Exception):
+    pass
+
+
+def _fix(argument):
+    def inner(body):
+        recursion_limit = [20]
+
+        def fix(x):
+            def r(z):
+                recursion_limit[0] -= 1
+                if recursion_limit[0] <= 0:
+                    raise RecursionDepthExceeded()
+                else:
+                    return fix(z)
+
+            return body(r)(x)
+        return fix(argument)
+
+    return inner
+
+
+def curry(f): return lambda x: lambda y: f((x, y))
+
+
+def _fix2(a1):
+    return lambda a2: lambda body: \
+        _fix((a1, a2))(lambda r: lambda n_l: body(curry(r))(n_l[0])(n_l[1]))
+
+
+primitiveRecursion1 = Primitive("fix1",
+                                arrow(t0,
+                                      arrow(arrow(t0, t1), t0, t1),
+                                      t1),
+                                _fix)
+
+primitiveRecursion2 = Primitive("fix2",
+                                arrow(t0, t1,
+                                      arrow(arrow(t0, t1, t2), t0, t1, t2),
+                                      t2),
+                                _fix2)
+
+
+def _match(l):
+    return lambda b: lambda f: b if l == [] else f(l[0])(l[1:])
+
+
+def primitives():
+    return [Primitive(str(j), tint, j) for j in range(6)] + [
+        Primitive("empty", tlist(t0), []),
+        Primitive("singleton", arrow(t0, tlist(t0)), _single),
+        Primitive("range", arrow(tint, tlist(tint)), _range),
+        Primitive("++", arrow(tlist(t0), tlist(t0), tlist(t0)), _append),
+        # Primitive("map", arrow(arrow(t0, t1), tlist(t0), tlist(t1)), _map),
+        Primitive(
+            "mapi",
+            arrow(
+                arrow(
+                    tint,
+                    t0,
+                    t1),
+                tlist(t0),
+                tlist(t1)),
+            _mapi),
+        # Primitive("reduce", arrow(arrow(t1, t0, t1), t1, tlist(t0), t1), _reduce),
+        Primitive(
+            "reducei",
+            arrow(
+                arrow(
+                    tint,
+                    t1,
+                    t0,
+                    t1),
+                t1,
+                tlist(t0),
+                t1),
+            _reducei),
+
+        Primitive("true", tbool, True),
+        Primitive("not", arrow(tbool, tbool), _not),
+        Primitive("and", arrow(tbool, tbool, tbool), _and),
+        Primitive("or", arrow(tbool, tbool, tbool), _or),
+        # Primitive("if", arrow(tbool, t0, t0, t0), _if),
+
+        Primitive("sort", arrow(tlist(tint), tlist(tint)), sorted),
+        Primitive("+", arrow(tint, tint, tint), _addition),
+        Primitive("*", arrow(tint, tint, tint), _multiplication),
+        Primitive("negate", arrow(tint, tint), _negate),
+        Primitive("mod", arrow(tint, tint, tint), _mod),
+        Primitive("eq?", arrow(tint, tint, tbool), _eq),
+        Primitive("gt?", arrow(tint, tint, tbool), _gt),
+        Primitive("is-prime", arrow(tint, tbool), _isPrime),
+        Primitive("is-square", arrow(tint, tbool), _isSquare),
+
+        # these are achievable with above primitives, but unlikely
+        #Primitive("flatten", arrow(tlist(tlist(t0)), tlist(t0)), _flatten),
+        # (lambda (reduce (lambda (lambda (++ $1 $0))) empty $0))
+        Primitive("sum", arrow(tlist(tint), tint), sum),
+        # (lambda (lambda (reduce (lambda (lambda (+ $0 $1))) 0 $0)))
+        Primitive("reverse", arrow(tlist(t0), tlist(t0)), _reverse),
+        # (lambda (reduce (lambda (lambda (++ (singleton $0) $1))) empty $0))
+        Primitive("all", arrow(arrow(t0, tbool), tlist(t0), tbool), _all),
+        # (lambda (lambda (reduce (lambda (lambda (and $0 $1))) true (map $1 $0))))
+        Primitive("any", arrow(arrow(t0, tbool), tlist(t0), tbool), _any),
+        # (lambda (lambda (reduce (lambda (lambda (or $0 $1))) true (map $1 $0))))
+        Primitive("index", arrow(tint, tlist(t0), t0), _index),
+        # (lambda (lambda (reducei (lambda (lambda (lambda (if (eq? $1 $4) $0 0)))) 0 $0)))
+        Primitive("filter", arrow(arrow(t0, tbool), tlist(t0), tlist(t0)), _filter),
+        # (lambda (lambda (reduce (lambda (lambda (++ $1 (if ($3 $0) (singleton $0) empty)))) empty $0)))
+        #Primitive("replace", arrow(arrow(tint, t0, tbool), tlist(t0), tlist(t0), tlist(t0)), _replace),
+        # (FLATTEN (lambda (lambda (lambda (mapi (lambda (lambda (if ($4 $1 $0) $3 (singleton $1)))) $0)))))
+        Primitive("slice", arrow(tint, tint, tlist(t0), tlist(t0)), _slice),
+        # (lambda (lambda (lambda (reducei (lambda (lambda (lambda (++ $2 (if (and (or (gt? $1 $5) (eq? $1 $5)) (not (or (gt? $4 $1) (eq? $1 $4)))) (singleton $0) empty))))) empty $0))))
+    ]
+
+
+def basePrimitives():
+    return [Primitive(str(j), tint, j) for j in range(6)] + [
+        Primitive("*", arrow(tint, tint, tint), _multiplication),
+        Primitive("gt?", arrow(tint, tint, tbool), _gt),
+        Primitive("is-prime", arrow(tint, tbool), _isPrime),
+        Primitive("is-square", arrow(tint, tbool), _isSquare),
+        # McCarthy
+        Primitive("empty", tlist(t0), []),
+        Primitive("cons", arrow(t0, tlist(t0), tlist(t0)), _cons),
+        Primitive("car", arrow(tlist(t0), t0), _car),
+        Primitive("cdr", arrow(tlist(t0), tlist(t0)), _cdr),
+        Primitive("empty?", arrow(tlist(t0), tbool), _isEmpty),
+        Primitive("if", arrow(tbool, t0, t0, t0), _if),
+        Primitive("eq?", arrow(tint, tint, tbool), _eq),
+        Primitive("+", arrow(tint, tint, tint), _addition),
+        Primitive("-", arrow(tint, tint, tint), _subtraction)
+    ]
+
+zip_primitive = Primitive("zip", arrow(tlist(t0), tlist(t1), arrow(t0, t1, t2), tlist(t2)), _zip)
+
+def bootstrapTarget():
+    """These are the primitives that we hope to learn from the bootstrapping procedure"""
+    return [
+        # learned primitives
+        Primitive("map", arrow(arrow(t0, t1), tlist(t0), tlist(t1)), _map),
+        Primitive("unfold", arrow(t0, arrow(t0,tbool), arrow(t0,t1), arrow(t0,t0), tlist(t1)), _unfold),
+        Primitive("range", arrow(tint, tlist(tint)), _range),
+        Primitive("index", arrow(tint, tlist(t0), t0), _index),
+        Primitive("fold", arrow(tlist(t0), t1, arrow(t0, t1, t1), t1), _fold),
+        Primitive("length", arrow(tlist(t0), tint), len),
+
+        # built-ins
+        Primitive("if", arrow(tbool, t0, t0, t0), _if),
+        Primitive("+", arrow(tint, tint, tint), _addition),
+        Primitive("-", arrow(tint, tint, tint), _subtraction),
+        Primitive("empty", tlist(t0), []),
+        Primitive("cons", arrow(t0, tlist(t0), tlist(t0)), _cons),
+        Primitive("car", arrow(tlist(t0), t0), _car),
+        Primitive("cdr", arrow(tlist(t0), tlist(t0)), _cdr),
+        Primitive("empty?", arrow(tlist(t0), tbool), _isEmpty),
+    ] + [Primitive(str(j), tint, j) for j in range(2)]
+
+
+def bootstrapTarget_extra():
+    """This is the bootstrap target plus list domain specific stuff"""
+    return bootstrapTarget() + [
+        Primitive("*", arrow(tint, tint, tint), _multiplication),
+        Primitive("mod", arrow(tint, tint, tint), _mod),
+        Primitive("gt?", arrow(tint, tint, tbool), _gt),
+        Primitive("eq?", arrow(tint, tint, tbool), _eq),
+        Primitive("is-prime", arrow(tint, tbool), _isPrime),
+        Primitive("is-square", arrow(tint, tbool), _isSquare),
+    ]
+
+def no_length():
+    """this is the primitives without length because one of the reviewers wanted this"""
+    return [p for p in bootstrapTarget() if p.name != "length"] + [
+        Primitive("*", arrow(tint, tint, tint), _multiplication),
+        Primitive("mod", arrow(tint, tint, tint), _mod),
+        Primitive("gt?", arrow(tint, tint, tbool), _gt),
+        Primitive("eq?", arrow(tint, tint, tbool), _eq),
+        Primitive("is-prime", arrow(tint, tbool), _isPrime),
+        Primitive("is-square", arrow(tint, tbool), _isSquare),
+    ]
+
+
+def McCarthyPrimitives():
+    "These are < primitives provided by 1959 lisp as introduced by McCarthy"
+    return [
+        Primitive("empty", tlist(t0), []),
+        Primitive("cons", arrow(t0, tlist(t0), tlist(t0)), _cons),
+        Primitive("car", arrow(tlist(t0), t0), _car),
+        Primitive("cdr", arrow(tlist(t0), tlist(t0)), _cdr),
+        Primitive("empty?", arrow(tlist(t0), tbool), _isEmpty),
+        #Primitive("unfold", arrow(t0, arrow(t0,t1), arrow(t0,t0), arrow(t0,tbool), tlist(t1)), _isEmpty),
+        #Primitive("1+", arrow(tint,tint),None),
+        # Primitive("range", arrow(tint, tlist(tint)), range),
+        # Primitive("map", arrow(arrow(t0, t1), tlist(t0), tlist(t1)), _map),
+        # Primitive("index", arrow(tint,tlist(t0),t0),None),
+        # Primitive("length", arrow(tlist(t0),tint),None),
+        primitiveRecursion1,
+        #primitiveRecursion2,
+        Primitive("gt?", arrow(tint, tint, tbool), _gt),
+        Primitive("if", arrow(tbool, t0, t0, t0), _if),
+        Primitive("eq?", arrow(tint, tint, tbool), _eq),
+        Primitive("+", arrow(tint, tint, tint), _addition),
+        Primitive("-", arrow(tint, tint, tint), _subtraction),
+    ] + [Primitive(str(j), tint, j) for j in range(2)]
+
+
+if __name__ == "__main__":
+    bootstrapTarget()
+    g = Grammar.uniform(McCarthyPrimitives())
+    # with open("/home/ellisk/om/ec/experimentOutputs/list_aic=1.0_arity=3_ET=1800_expandFrontier=2.0_it=4_likelihoodModel=all-or-nothing_MF=5_baseline=False_pc=10.0_L=1.0_K=5_rec=False.pickle", "rb") as handle:
+    #     b = pickle.load(handle).grammars[-1]
+    # print b
+
+    p = Program.parse(
+        "(lambda (lambda (lambda (if (empty? $0) empty (cons (+ (car $1) (car $0)) ($2 (cdr $1) (cdr $0)))))))")
+    t = arrow(tlist(tint), tlist(tint), tlist(tint))  # ,tlist(tbool))
+    print(g.logLikelihood(arrow(t, t), p))
+    assert False
+    print(b.logLikelihood(arrow(t, t), p))
+
+    # p = Program.parse("""(lambda (lambda
+    # (unfold 0
+    # (lambda (+ (index $0 $2) (index $0 $1)))
+    # (lambda (1+ $0))
+    # (lambda (eq? $0 (length $1))))))
+    # """)
+    p = Program.parse("""(lambda (lambda
+    (map (lambda (+ (index $0 $2) (index $0 $1))) (range (length $0))  )))""")
+    # .replace("unfold", "#(lambda (lambda (lambda (lambda (fix1 $0 (lambda (lambda (#(lambda (lambda (lambda (if $0 empty (cons $1 $2))))) ($1 ($3 $0)) ($4 $0) ($5 $0)))))))))").\
+    # replace("length", "#(lambda (fix1 $0 (lambda (lambda (if (empty? $0) 0 (+ ($1 (cdr $0)) 1))))))").\
+    # replace("forloop", "(#(lambda (lambda (lambda (lambda (fix1 $0 (lambda (lambda (#(lambda (lambda (lambda (if $0 empty (cons $1 $2))))) ($1 ($3 $0)) ($4 $0) ($5 $0))))))))) (lambda (#(eq? 0) $0)) $0 (lambda (#(lambda (- $0 1)) $0)))").\
+    # replace("inc","#(lambda (+ $0 1))").\
+    # replace("drop","#(lambda (lambda (fix2 $0 $1 (lambda (lambda (lambda (if
+    # (#(eq? 0) $1) $0 (cdr ($2 (- $1 1) $0)))))))))"))
+    print(p)
+    print(g.logLikelihood(t, p))
+    assert False
+
+    print("??")
+    p = Program.parse(
+        "#(lambda (#(lambda (lambda (lambda (fix1 $0 (lambda (lambda (if (empty? $0) $3 ($4 (car $0) ($1 (cdr $0)))))))))) (lambda $1) 1))")
+    for j in range(10):
+        l = list(range(j))
+        print(l, p.evaluate([])(lambda x: x * 2)(l))
+        print()
+    print()
+
+    print("multiply")
+    p = Program.parse(
+        "(lambda (lambda (lambda (if (eq? $0 0) 0 (+ $1 ($2 $1 (- $0 1)))))))")
+    print(g.logLikelihood(arrow(arrow(tint, tint, tint), tint, tint, tint), p))
+    print()
+
+    print("take until 0")
+    p = Program.parse("(lambda (lambda (if (eq? $1 0) empty (cons $1 $0))))")
+    print(g.logLikelihood(arrow(tint, tlist(tint), tlist(tint)), p))
+    print()
+
+    print("countdown primitive")
+    p = Program.parse(
+        "(lambda (lambda (if (eq? $0 0) empty (cons (+ $0 1) ($1 (- $0 1))))))")
+    print(
+        g.logLikelihood(
+            arrow(
+                arrow(
+                    tint, tlist(tint)), arrow(
+                    tint, tlist(tint))), p))
+    print(_fix(9)(p.evaluate([])))
+    print("countdown w/ better primitives")
+    p = Program.parse(
+        "(lambda (lambda (if (eq0 $0) empty (cons (+1 $0) ($1 (-1 $0))))))")
+    print(
+        g.logLikelihood(
+            arrow(
+                arrow(
+                    tint, tlist(tint)), arrow(
+                    tint, tlist(tint))), p))
+
+    print()
+
+    print("prepend zeros")
+    p = Program.parse(
+        "(lambda (lambda (lambda (if (eq? $1 0) $0 (cons 0 ($2 (- $1 1) $0))))))")
+    print(
+        g.logLikelihood(
+            arrow(
+                arrow(
+                    tint,
+                    tlist(tint),
+                    tlist(tint)),
+                tint,
+                tlist(tint),
+                tlist(tint)),
+            p))
+    print()
+    assert False
+
+    p = Program.parse(
+        "(lambda (fix1 $0 (lambda (lambda (if (empty? $0) 0 (+ 1 ($1 (cdr $0))))))))")
+    print(p.evaluate([])(list(range(17))))
+    print(g.logLikelihood(arrow(tlist(tbool), tint), p))
+
+    p = Program.parse(
+        "(lambda (lambda (if (empty? $0) 0 (+ 1 ($1 (cdr $0))))))")
+    print(
+        g.logLikelihood(
+            arrow(
+                arrow(
+                    tlist(tbool), tint), arrow(
+                    tlist(tbool), tint)), p))
+
+    p = Program.parse(
+        "(lambda (fix1 $0 (lambda (lambda (if (empty? $0) 0 (+ (car $0) ($1 (cdr $0))))))))")
+
+    print(p.evaluate([])(list(range(4))))
+    print(g.logLikelihood(arrow(tlist(tint), tint), p))
+
+    p = Program.parse(
+        "(lambda (lambda (if (empty? $0) 0 (+ (car $0) ($1 (cdr $0))))))")
+    print(p)
+    print(
+        g.logLikelihood(
+            arrow(
+                arrow(
+                    tlist(tint),
+                    tint),
+                tlist(tint),
+                tint),
+            p))
+
+    print("take")
+    p = Program.parse(
+        "(lambda (lambda (lambda (if (eq? $1 0) empty (cons (car $0) ($2 (- $1 1) (cdr $0)))))))")
+    print(p)
+    print(
+        g.logLikelihood(
+            arrow(
+                arrow(
+                    tint,
+                    tlist(tint),
+                    tlist(tint)),
+                tint,
+                tlist(tint),
+                tlist(tint)),
+            p))
+    assert False
+
+    print(p.evaluate([])(list(range(4))))
+    print(g.logLikelihood(arrow(tlist(tint), tlist(tint)), p))
+
+    p = Program.parse(
+        """(lambda (fix (lambda (lambda (match $0 0 (lambda (lambda (+ $1 ($3 $0))))))) $0))""")
+    print(p.evaluate([])(list(range(4))))
+    print(g.logLikelihood(arrow(tlist(tint), tint), p))

dreamcoder/domains/list/main.py

@@ -0,0 +1,410 @@
+import random
+from collections import defaultdict
+import json
+import math
+import os
+import datetime
+
+from dreamcoder.dreamcoder import explorationCompression
+from dreamcoder.utilities import eprint, flatten, testTrainSplit
+from dreamcoder.grammar import Grammar
+from dreamcoder.task import Task
+from dreamcoder.type import Context, arrow, tbool, tlist, tint, t0, UnificationFailure
+from dreamcoder.domains.list.listPrimitives import basePrimitives, primitives, McCarthyPrimitives, bootstrapTarget_extra, no_length
+from dreamcoder.domains.list.makeListTasks import make_list_bootstrap_tasks, sortBootstrap, EASYLISTTASKS
+
+
+def retrieveJSONTasks(filename, features=False):
+    """
+    For JSON of the form:
+        {"name": str,
+         "type": {"input" : bool|int|list-of-bool|list-of-int,
+                  "output": bool|int|list-of-bool|list-of-int},
+         "examples": [{"i": data, "o": data}]}
+    """
+    with open(filename, "r") as f:
+        loaded = json.load(f)
+    TP = {
+        "bool": tbool,
+        "int": tint,
+        "list-of-bool": tlist(tbool),
+        "list-of-int": tlist(tint),
+    }
+    return [Task(
+        item["name"],
+        arrow(TP[item["type"]["input"]], TP[item["type"]["output"]]),
+        [((ex["i"],), ex["o"]) for ex in item["examples"]],
+        features=(None if not features else list_features(
+            [((ex["i"],), ex["o"]) for ex in item["examples"]])),
+        cache=False,
+    ) for item in loaded]
+
+
+def list_features(examples):
+    if any(isinstance(i, int) for (i,), _ in examples):
+        # obtain features for number inputs as list of numbers
+        examples = [(([i],), o) for (i,), o in examples]
+    elif any(not isinstance(i, list) for (i,), _ in examples):
+        # can't handle non-lists
+        return []
+    elif any(isinstance(x, list) for (xs,), _ in examples for x in xs):
+        # nested lists are hard to extract features for, so we'll
+        # obtain features as if flattened
+        examples = [(([x for xs in ys for x in xs],), o)
+                    for (ys,), o in examples]
+
+    # assume all tasks have the same number of examples
+    # and all inputs are lists
+    features = []
+    ot = type(examples[0][1])
+
+    def mean(l): return 0 if not l else sum(l) / len(l)
+    imean = [mean(i) for (i,), o in examples]
+    ivar = [sum((v - imean[idx])**2
+                for v in examples[idx][0][0])
+            for idx in range(len(examples))]
+
+    # DISABLED length of each input and output
+    # total difference between length of input and output
+    # DISABLED normalized count of numbers in input but not in output
+    # total normalized count of numbers in input but not in output
+    # total difference between means of input and output
+    # total difference between variances of input and output
+    # output type (-1=bool, 0=int, 1=list)
+    # DISABLED outputs if integers, else -1s
+    # DISABLED outputs if bools (-1/1), else 0s
+    if ot == list:  # lists of ints or bools
+        omean = [mean(o) for (i,), o in examples]
+        ovar = [sum((v - omean[idx])**2
+                    for v in examples[idx][1])
+                for idx in range(len(examples))]
+
+        def cntr(
+            l, o): return 0 if not l else len(
+            set(l).difference(
+                set(o))) / len(l)
+        cnt_not_in_output = [cntr(i, o) for (i,), o in examples]
+
+        #features += [len(i) for (i,), o in examples]
+        #features += [len(o) for (i,), o in examples]
+        features.append(sum(len(i) - len(o) for (i,), o in examples))
+        #features += cnt_not_int_output
+        features.append(sum(cnt_not_in_output))
+        features.append(sum(om - im for im, om in zip(imean, omean)))
+        features.append(sum(ov - iv for iv, ov in zip(ivar, ovar)))
+        features.append(1)
+        # features += [-1 for _ in examples]
+        # features += [0 for _ in examples]
+    elif ot == bool:
+        outs = [o for (i,), o in examples]
+
+        #features += [len(i) for (i,), o in examples]
+        #features += [-1 for _ in examples]
+        features.append(sum(len(i) for (i,), o in examples))
+        #features += [0 for _ in examples]
+        features.append(0)
+        features.append(sum(imean))
+        features.append(sum(ivar))
+        features.append(-1)
+        # features += [-1 for _ in examples]
+        # features += [1 if o else -1 for o in outs]
+    else:  # int
+        def cntr(
+            l, o): return 0 if not l else len(
+            set(l).difference(
+                set(o))) / len(l)
+        cnt_not_in_output = [cntr(i, [o]) for (i,), o in examples]
+        outs = [o for (i,), o in examples]
+
+        #features += [len(i) for (i,), o in examples]
+        #features += [1 for (i,), o in examples]
+        features.append(sum(len(i) for (i,), o in examples))
+        #features += cnt_not_int_output
+        features.append(sum(cnt_not_in_output))
+        features.append(sum(o - im for im, o in zip(imean, outs)))
+        features.append(sum(ivar))
+        features.append(0)
+        # features += outs
+        # features += [0 for _ in examples]
+
+    return features
+
+
+def isListFunction(tp):
+    try:
+        Context().unify(tp, arrow(tlist(tint), t0))
+        return True
+    except UnificationFailure:
+        return False
+
+
+def isIntFunction(tp):
+    try:
+        Context().unify(tp, arrow(tint, t0))
+        return True
+    except UnificationFailure:
+        return False
+
+try:
+    from dreamcoder.recognition import RecurrentFeatureExtractor
+    class LearnedFeatureExtractor(RecurrentFeatureExtractor):
+        H = 64
+
+        special = None
+
+        def tokenize(self, examples):
+            def sanitize(l): return [z if z in self.lexicon else "?"
+                                     for z_ in l
+                                     for z in (z_ if isinstance(z_, list) else [z_])]
+
+            tokenized = []
+            for xs, y in examples:
+                if isinstance(y, list):
+                    y = ["LIST_START"] + y + ["LIST_END"]
+                else:
+                    y = [y]
+                y = sanitize(y)
+                if len(y) > self.maximumLength:
+                    return None
+
+                serializedInputs = []
+                for xi, x in enumerate(xs):
+                    if isinstance(x, list):
+                        x = ["LIST_START"] + x + ["LIST_END"]
+                    else:
+                        x = [x]
+                    x = sanitize(x)
+                    if len(x) > self.maximumLength:
+                        return None
+                    serializedInputs.append(x)
+
+                tokenized.append((tuple(serializedInputs), y))
+
+            return tokenized
+
+        def __init__(self, tasks, testingTasks=[], cuda=False):
+            self.lexicon = set(flatten((t.examples for t in tasks + testingTasks), abort=lambda x: isinstance(
+                x, str))).union({"LIST_START", "LIST_END", "?"})
+
+            # Calculate the maximum length
+            self.maximumLength = float('inf') # Believe it or not this is actually important to have here
+            self.maximumLength = max(len(l)
+                                     for t in tasks + testingTasks
+                                     for xs, y in self.tokenize(t.examples)
+                                     for l in [y] + [x for x in xs])
+
+            self.recomputeTasks = True
+
+            super(
+                LearnedFeatureExtractor,
+                self).__init__(
+                lexicon=list(
+                    self.lexicon),
+                tasks=tasks,
+                cuda=cuda,
+                H=self.H,
+                bidirectional=True)
+except: pass
+
+def train_necessary(t):
+    if t.name in {"head", "is-primes", "len", "pop", "repeat-many", "tail", "keep primes", "keep squares"}:
+        return True
+    if any(t.name.startswith(x) for x in {
+        "add-k", "append-k", "bool-identify-geq-k", "count-k", "drop-k",
+        "empty", "evens", "has-k", "index-k", "is-mod-k", "kth-largest",
+        "kth-smallest", "modulo-k", "mult-k", "remove-index-k",
+        "remove-mod-k", "repeat-k", "replace-all-with-index-k", "rotate-k",
+        "slice-k-n", "take-k",
+    }):
+        return "some"
+    return False
+
+
+def list_options(parser):
+    parser.add_argument(
+        "--noMap", action="store_true", default=False,
+        help="Disable built-in map primitive")
+    parser.add_argument(
+        "--noUnfold", action="store_true", default=False,
+        help="Disable built-in unfold primitive")
+    parser.add_argument(
+        "--noLength", action="store_true", default=False,
+        help="Disable built-in length primitive")
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="Lucas-old",
+        choices=[
+            "bootstrap",
+            "sorting",
+            "Lucas-old",
+            "Lucas-depth1",
+            "Lucas-depth2",
+            "Lucas-depth3"])
+    parser.add_argument("--maxTasks", type=int,
+                        default=None,
+                        help="truncate tasks to fit within this boundary")
+    parser.add_argument("--primitives",
+                        default="common",
+                        help="Which primitive set to use",
+                        choices=["McCarthy", "base", "rich", "common", "noLength"])
+    parser.add_argument("--extractor", type=str,
+                        choices=["hand", "deep", "learned"],
+                        default="learned")
+    parser.add_argument("--split", metavar="TRAIN_RATIO",
+                        type=float,
+                        help="split test/train")
+    parser.add_argument("-H", "--hidden", type=int,
+                        default=64,
+                        help="number of hidden units")
+    parser.add_argument("--random-seed", type=int, default=17)
+
+
+def main(args):
+    """
+    Takes the return value of the `commandlineArguments()` function as input and
+    trains/tests the model on manipulating sequences of numbers.
+    """
+    random.seed(args.pop("random_seed"))
+
+    dataset = args.pop("dataset")
+    tasks = {
+        "Lucas-old": lambda: retrieveJSONTasks("data/list_tasks.json") + sortBootstrap(),
+        "bootstrap": make_list_bootstrap_tasks,
+        "sorting": sortBootstrap,
+        "Lucas-depth1": lambda: retrieveJSONTasks("data/list_tasks2.json")[:105],
+        "Lucas-depth2": lambda: retrieveJSONTasks("data/list_tasks2.json")[:4928],
+        "Lucas-depth3": lambda: retrieveJSONTasks("data/list_tasks2.json"),
+    }[dataset]()
+
+    maxTasks = args.pop("maxTasks")
+    if maxTasks and len(tasks) > maxTasks:
+        necessaryTasks = []  # maxTasks will not consider these
+        if dataset.startswith("Lucas2.0") and dataset != "Lucas2.0-depth1":
+            necessaryTasks = tasks[:105]
+
+        eprint("Unwilling to handle {} tasks, truncating..".format(len(tasks)))
+        random.shuffle(tasks)
+        del tasks[maxTasks:]
+        tasks = necessaryTasks + tasks
+
+    if dataset.startswith("Lucas"):
+        # extra tasks for filter
+        tasks.extend([
+            Task("remove empty lists",
+                 arrow(tlist(tlist(tbool)), tlist(tlist(tbool))),
+                 [((ls,), list(filter(lambda l: len(l) > 0, ls)))
+                  for _ in range(15)
+                  for ls in [[[random.random() < 0.5 for _ in range(random.randint(0, 3))]
+                              for _ in range(4)]]]),
+            Task("keep squares",
+                 arrow(tlist(tint), tlist(tint)),
+                 [((xs,), list(filter(lambda x: int(math.sqrt(x)) ** 2 == x,
+                                      xs)))
+                  for _ in range(15)
+                  for xs in [[random.choice([0, 1, 4, 9, 16, 25])
+                              if random.random() < 0.5
+                              else random.randint(0, 9)
+                              for _ in range(7)]]]),
+            Task("keep primes",
+                 arrow(tlist(tint), tlist(tint)),
+                 [((xs,), list(filter(lambda x: x in {2, 3, 5, 7, 11, 13, 17,
+                                                      19, 23, 29, 31, 37}, xs)))
+                  for _ in range(15)
+                  for xs in [[random.choice([2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37])
+                              if random.random() < 0.5
+                              else random.randint(0, 9)
+                              for _ in range(7)]]]),
+        ])
+        for i in range(4):
+            tasks.extend([
+                Task("keep eq %s" % i,
+                     arrow(tlist(tint), tlist(tint)),
+                     [((xs,), list(filter(lambda x: x == i, xs)))
+                      for _ in range(15)
+                      for xs in [[random.randint(0, 6) for _ in range(5)]]]),
+                Task("remove eq %s" % i,
+                     arrow(tlist(tint), tlist(tint)),
+                     [((xs,), list(filter(lambda x: x != i, xs)))
+                      for _ in range(15)
+                      for xs in [[random.randint(0, 6) for _ in range(5)]]]),
+                Task("keep gt %s" % i,
+                     arrow(tlist(tint), tlist(tint)),
+                     [((xs,), list(filter(lambda x: x > i, xs)))
+                      for _ in range(15)
+                      for xs in [[random.randint(0, 6) for _ in range(5)]]]),
+                Task("remove gt %s" % i,
+                     arrow(tlist(tint), tlist(tint)),
+                     [((xs,), list(filter(lambda x: not x > i, xs)))
+                      for _ in range(15)
+                      for xs in [[random.randint(0, 6) for _ in range(5)]]])
+            ])
+
+    def isIdentityTask(t):
+        return all( len(xs) == 1 and xs[0] == y for xs, y in t.examples  )
+    eprint("Removed", sum(isIdentityTask(t) for t in tasks), "tasks that were just the identity function")
+    tasks = [t for t in tasks if not isIdentityTask(t) ]
+
+    prims = {"base": basePrimitives,
+             "McCarthy": McCarthyPrimitives,
+             "common": bootstrapTarget_extra,
+             "noLength": no_length,
+             "rich": primitives}[args.pop("primitives")]()
+    haveLength = not args.pop("noLength")
+    haveMap = not args.pop("noMap")
+    haveUnfold = not args.pop("noUnfold")
+    eprint(f"Including map as a primitive? {haveMap}")
+    eprint(f"Including length as a primitive? {haveLength}")
+    eprint(f"Including unfold as a primitive? {haveUnfold}")
+    baseGrammar = Grammar.uniform([p
+                                   for p in prims
+                                   if (p.name != "map" or haveMap) and \
+                                   (p.name != "unfold" or haveUnfold) and \
+                                   (p.name != "length" or haveLength)])
+
+    extractor = {
+        "learned": LearnedFeatureExtractor,
+    }[args.pop("extractor")]
+    extractor.H = args.pop("hidden")
+
+    timestamp = datetime.datetime.now().isoformat()
+    outputDirectory = "experimentOutputs/list/%s"%timestamp
+    os.system("mkdir -p %s"%outputDirectory)
+    
+    args.update({
+        "featureExtractor": extractor,
+        "outputPrefix": "%s/list"%outputDirectory,
+        "evaluationTimeout": 0.0005,
+    })
+    
+
+    eprint("Got {} list tasks".format(len(tasks)))
+    split = args.pop("split")
+    if split:
+        train_some = defaultdict(list)
+        for t in tasks:
+            necessary = train_necessary(t)
+            if not necessary:
+                continue
+            if necessary == "some":
+                train_some[t.name.split()[0]].append(t)
+            else:
+                t.mustTrain = True
+        for k in sorted(train_some):
+            ts = train_some[k]
+            random.shuffle(ts)
+            ts.pop().mustTrain = True
+
+        test, train = testTrainSplit(tasks, split)
+        if True:
+            test = [t for t in test
+                    if t.name not in EASYLISTTASKS]
+
+        eprint(
+            "Alotted {} tasks for training and {} for testing".format(
+                len(train), len(test)))
+    else:
+        train = tasks
+        test = []
+
+    explorationCompression(baseGrammar, train, testingTasks=test, **args)

dreamcoder/domains/list/makeListTasks.py

@@ -0,0 +1,587 @@
+
+
+from dreamcoder.type import *
+from dreamcoder.task import Task
+from dreamcoder.utilities import eprint, hashable
+
+from random import randint, random, seed
+from itertools import product
+
+# Excluded routines either impossible or astronomically improbable
+# I'm cutting these off at ~20 nats in learned grammars.
+EXCLUDES = {
+    "dedup",
+    "intersperse-k",
+    "pow-base-k",
+    "prime",
+    "replace-all-k-with-n",
+    "replace-index-k-with-n",
+    "uniq",
+}
+
+# These are tasks that are easy (solved from base DSL) and also uninteresting
+# We exclude these from the test set
+EASYLISTTASKS = {
+    "add-k with k=2",
+    "bool-identify-geq-k with k=2",
+    "bool-identify-geq-k with k=3",
+    "bool-identify-is-mod-k with k=1",
+    "bool-identify-is-prime",
+    "bool-identify-k with k=0",
+    "bool-identify-k with k=1",
+    "bool-identify-k with k=2",
+    "caesar-cipher-k-modulo-n with k=3 and n=2",
+    "drop-k with k=1",
+    "drop-k with k=2",
+    "drop-k with k=4",
+    "index-head",
+    "index-k with k=2",
+    "index-k with k=4",
+    "is-mod-k with k=1",
+    "is-odds",
+    "is-squares",
+    "pow-k with k=2",
+    "pow-k with k=3",
+    "prepend-index-k with k=3",
+    "prepend-index-k with k=5",
+    "prepend-k with k=1",
+    "prepend-k with k=2",
+    "prepend-k with k=3",
+    "remove-index-k with k=1",
+    "replace-all-with-index-k with k=2",
+    "replace-all-with-index-k with k=3",
+    "slice-k-n with k=1 and n=2",
+    "slice-k-n with k=2 and n=1",
+    "slice-k-n with k=3 and n=1",
+}
+
+def make_list_task(name, examples, **params):
+    input_type = guess_type([i for (i,), _ in examples])
+    output_type = guess_type([o for _, o in examples])
+
+    # We can internally handle lists of bools.
+    # We explicitly create these by modifying existing routines.
+    if name.startswith("identify"):
+        boolexamples = [((i,), list(map(bool, o))) for (i,), o in examples]
+        yield from make_list_task("bool-" + name, boolexamples, **params)
+        # for now, we'll stick with the boolean-only tasks and not have a copy
+        # for integers.
+        return
+
+    program_type = arrow(input_type, output_type)
+    cache = all(hashable(x) for x in examples)
+
+    if params:
+        eq_params = ["{}={}".format(k, v) for k, v in params.items()]
+        if len(eq_params) == 1:
+            ext = eq_params[0]
+        elif len(eq_params) == 2:
+            ext = "{} and {}".format(*eq_params)
+        else:
+            ext = ", ".join(eq_params[:-1])
+            ext = "{}, and {}".format(ext, eq_params[-1])
+        name += " with " + ext
+
+    yield Task(name, program_type, examples, cache=cache)
+
+
+def make_list_tasks(n_examples):
+    import listroutines as lr
+
+    for routine in lr.find(count=100):  # all routines
+        if routine.id in EXCLUDES:
+            continue
+        if routine.is_parametric():
+            keys = list(routine.example_params()[0].keys())
+            for params in map(lambda values: dict(zip(keys, values)),
+                              product(range(6), repeat=len(keys))):
+                try:
+                    if routine.id == "rotate-k":
+                        # rotate-k is hard if list is smaller than k
+                        k = params["k"]
+                        if k < 1:
+                            continue
+                        inps = []
+                        for _ in range(n_examples):
+                            r = randint(abs(k) + 1, 17)
+                            inp = routine.gen(len=r, **params)[0]
+                            inps.append(inp)
+                    else:
+                        inps = routine.gen(count=n_examples, **params)
+                    examples = [((inp,), routine.eval(inp, **params))
+                                for inp in inps]
+                    yield from make_list_task(routine.id, examples, **params)
+                except lr.APIError:  # invalid params
+                    continue
+        else:
+            inps = routine.examples()
+            if len(inps) > n_examples:
+                inps = inps[:n_examples]
+            elif len(inps) < n_examples:
+                inps += routine.gen(count=(n_examples - len(inps)))
+            examples = [((inp,), routine.eval(inp)) for inp in inps]
+            yield from make_list_task(routine.id, examples)
+
+
+def make_list_bootstrap_tasks():
+    seed(42)
+
+    def suffixes(l):
+        if l == []:
+            return []
+        else:
+            return [l[1:]] + suffixes(l[1:])
+
+    def flip(): return random() > 0.5
+
+    def randomSuffix():
+        return [randint(0, 9) for _ in range(randint(1, 4))]
+
+    def randomList(minimum=0, minimumLength=4, maximumLength=6):
+        return [randint(minimum, 9) for _ in range(randint(minimumLength, maximumLength))]
+
+    def randomListOfLists():
+        return [randomSuffix() for _ in range(randint(2, 4))]
+
+    def randomListOfLists_bool(l=None):
+        if l is None:
+            l = randint(4, 7)
+        return [randomBooleanList() for _ in range(l)]
+
+    def randomBooleanList():
+        return [flip() for _ in range(randint(4, 7))]
+
+    # Reliably learned in under a minute; always triggers learning of length
+    # primitive
+    lengthBootstrap = [
+        # Task("length bool", arrow(tlist(tbool), tint),
+        #      [((l,), len(l))
+        #       for _ in range(10)
+        #       for l in [[flip() for _ in range(randint(0, 10))]]]),
+        Task("length int", arrow(tlist(tint), tint),
+             [((l,), len(l))
+              for _ in range(10)
+              for l in [randomList()]]),
+        Task("map length", arrow(tlist(tlist(tint)), tlist(tint)),
+             [((xss,), [len(xs) for xs in xss])
+              for _ in range(10)
+              for xss in [randomListOfLists()] ])
+    ]
+
+    # Encourages learning of unfolding
+    unfoldBootstrap = [
+        Task("countdown", arrow(tint, tlist(tint)),
+             [((n,), list(range(n + 1, 1, -1)))
+              for n in range(10)]),
+        Task("weird count", arrow(tint, tlist(tint)),
+             [((n,), list(range(-n,0,-1)))
+              for n in range(-10,0) ]),
+        Task("take every other", arrow(tlist(tint),tlist(tint)),
+             [((l,), [x for j,x in enumerate(l) if j%2 == 0])
+              for _ in range(9)
+              for l in [ [randint(0, 9) for _ in range(randint(1,4)*2)] ] ] + [(([],),[])]),
+        # Task("stutter every other", arrow(tlist(tint),tlist(tint)),
+        #      [((l,), [l[int(j/2)] for j in range(len(l)) ])
+        #       for _ in range(10)
+        #       for l in [ [randint(0, 9) for _ in range(randint(1,4)*2)] ] ]),
+        # Task("take until 3 reached", arrow(tlist(tint),tlist(tint)),
+        #      [((p + [3] + s,),p)
+        #       for _ in range(10)
+        #       for p in [ [z for z in randomList()[:5] if z != 3 ]]
+        #       for s in [randomList()] ]),
+        Task("drop last element", arrow(tlist(tint),tlist(tint)),
+             [((l,), l[:-1])
+              for _ in range(10)
+              for l in [ [randint(0, 9) for _ in range(randint(2,5))] ] ]),
+        # Task("suffixes", arrow(tlist(tint), tlist(tlist(tint))),
+        #      [((l,), suffixes(l))
+        #       for _ in range(10)
+        #       for l in [randomList()]]),
+        Task("range", arrow(tint, tlist(tint)),
+             [((n,), list(range(n)))
+              for n in range(10)]),
+        Task("range inclusive", arrow(tint, tlist(tint)),
+             [((n,), list(range(n + 1)))
+              for n in range(10)]),
+        # Task("range inclusive+1", arrow(tint, tlist(tint)),
+        #      [((n,), list(range(n + 2)))
+        #       for n in range(10)]),
+        # Task("range exclusive", arrow(tint, tlist(tint)),
+        #      [((n,), list(range(n - 1)))
+        #       for n in range(2, 11)]),
+        # Task("range length", arrow(tlist(tint),tlist(tint)),
+        #      [((l,),list(range(len(l))))
+        #       for _ in range(10)
+        #       for l in [randomList()] ])
+    ]
+
+    # Encourages learning how to treat a list as an array
+    arrayBootstrap = [
+        Task("index int", arrow(tint, tlist(tint), tint),
+             [((n, l), l[n])
+              for n in range(10)
+              for l in [[randint(0, 9) for _ in range(randint(n + 1, n + 5))]]]),
+        # Task("last n", arrow(tint, tlist(tint), tlist(tint)),
+        #      [((n, l), l[-n:])
+        #       for n in range(10)
+        #       for l in [[randint(0, 9) for _ in range(randint(n + 1, n + 5))]]]),
+        Task("1-index int", arrow(tint, tlist(tint), tint),
+             [((n, l), l[n - 1])
+              for n in range(1,11)
+              for l in [[randint(0, 9) for _ in range(randint(n + 1, n + 4))]]])
+        
+        # Task("index bool", arrow(tint, tlist(tbool), tbool),
+        #      [((n, l), l[n])
+        #       for n in range(10)
+        #       for l in [[flip() for _ in range(randint(n + 1, n + 5))]]])
+    ]
+
+    # Teaches how to slice lists, not sure if we really need this though
+    sliceBootstrap = [
+        Task("take bool", arrow(tint, tlist(tbool), tlist(tbool)),
+             [((n, l), l[:n])
+              for n in range(10)
+              for l in [[flip() for _ in range(randint(n, n + 5))]]]),
+        Task("drop bool", arrow(tint, tlist(tbool), tlist(tbool)),
+             [((n, l), l[n:])
+              for n in range(10)
+              for l in [[flip() for _ in range(randint(n, n + 5))]]]),
+
+        Task("take int", arrow(tint, tlist(tint), tlist(tint)),
+             [((n, l), l[:n])
+              for n in range(10)
+              for l in [[randint(0, 9) for _ in range(randint(n, n + 5))]]]),
+        Task("drop int", arrow(tint, tlist(tint), tlist(tint)),
+             [((n, l), l[n:])
+              for n in range(10)
+              for l in [[randint(0, 9) for _ in range(randint(n, n + 5))]]]),
+
+    ]
+
+    # learning to fold
+    foldBootstrap = [
+        Task("stutter", arrow(tlist(tint),tlist(tint)),
+             [((l,), [z for x in l for z in [x,x] ])
+              for _ in range(10)
+              for l in [randomList()] ]),
+        Task("sum", arrow(tlist(tint), tint),
+             [((l,), sum(l))
+              for _ in range(10)
+              for l in [randomList()]]),
+        # Task("difference", arrow(tlist(tint), tint),
+        #      [((l,), reduce(lambda x, y: y - x, reversed(l), 1))
+        #       for _ in range(10)
+        #       for l in [randomList()[:4]]]),
+        # Task("append bool", arrow(tlist(tbool), tlist(tbool), tlist(tbool)),
+        #      [((x, y), x + y)
+        #       for _ in range(10)
+        #       for [x, y] in [[randomBooleanList(), randomBooleanList()]]]),
+        Task("append constant 0", arrow(tlist(tint),tlist(tint)),
+             [((l,),l + [0])
+              for _ in range(10)
+              for l in [randomList()] ]),
+    ]
+
+    # learning to map
+    mapBootstrap = [
+        Task("map double", arrow(tlist(tint), tlist(tint)),
+             [((l,), list(map(lambda n: n * 2, l)))
+              for _ in range(10)
+              for l in [randomList()]]),
+        Task("map increment", arrow(tlist(tint),tlist(tint)),
+             [((l,),list(map(lambda n: n+1, l)))
+              for _ in range(10)
+              for l in [randomList()] ]),
+        Task("map negation", arrow(tlist(tint),tlist(tint)),
+             [((l,),list(map(lambda n: 0-n, l)))
+              for _ in range(10)
+              for l in [randomList()] ]),
+        # Task("map car", arrow(tlist(tlist(tint)), tlist(tint)),
+        #      [((l,), [n[0] for n in l])
+        #       for _ in4 range(10)
+        #       for l in [randomListOfLists()]]),
+        # Task("map cdr", arrow(tlist(tlist(tbool)),tlist(tlist(tbool))),
+        #      [((l,),map(lambda n: n[1:],l))
+        #       for _ in range(10)
+        #       for l in [randomListOfLists_bool()]]),
+        # Task("map empty?", arrow(tlist(tlist(tint)), tlist(tboolean)),
+        #      [((l,), [n == [] for n in l])
+        #       for _ in range(10)
+        #       for l in [[[] if flip() else randomList() for _ in range(randint(1, 5))]]]),
+
+        # Task("map eq 0?", arrow(tlist(tint),tlist(tboolean)),
+        #      [((l,),map(lambda n: 0 == n,l))
+        #       for _ in range(10)
+        #       for l in [[ randint(0,3) for _ in range(randint(4,7)) ]] ])
+
+    ]
+    difficultMaps = [
+                Task("map quadruple", arrow(tlist(tint), tlist(tint)),
+             [((l,), list(map(lambda n: n * 4, l)))
+              for _ in range(10)
+              for l in [randomList()]]),
+        Task("map add 3", arrow(tlist(tint),tlist(tint)),
+             [((l,),list(map(lambda n: n+3, l)))
+              for _ in range(10)
+              for l in [randomList()] ]),
+
+        ]
+
+    # Learning to zip lists together
+    zipBootstrap = [
+        Task("zip plus", arrow(tlist(tint),tlist(tint),tlist(tint)),
+             [((l1,l2),list(map(lambda x,y: x+y,l1,l2)))
+              for _ in range(10)
+              for l1 in [randomList(minimumLength=2, maximumLength=4)]
+              for l2 in [[ randint(0,9) for _ in range(len(l1)) ]]]),
+        Task("zip minus", arrow(tlist(tint),tlist(tint),tlist(tint)),
+             [((l1,l2),list(map(lambda x,y: x-y,l1,l2)))
+              for _ in range(10)
+              for l1 in [randomList(minimumLength=2, maximumLength=4)]
+              for l2 in [[ randint(0,9) for _ in range(len(l1)) ]]]),
+        # Task("zip eq?", arrow(tlist(tint), tlist(tint), tlist(tbool)),
+        #      [((l1, l2), list(map(lambda x, y: x == y, l1, l2)))
+        #       for _ in range(10)
+        #       for l1 in [[randint(0, 3) for _ in range(randint(4, 7))]]
+        #       for l2 in [[randint(0, 3) for _ in range(len(l1))]]]),
+        # Task("zip cons", arrow(tlist(tbool), tlist(tlist(tbool)), tlist(tlist(tbool))),
+        #      [((l1, l2), list(map(lambda x, y: [x] + y, l1, l2)))
+        #       for _ in range(10)
+        #       for l1 in [randomBooleanList()]
+        #       for l2 in [randomListOfLists_bool(l=len(l1))]]),
+        # Task("zip cons", arrow(tlist(tint),tlist(tlist(tint)),tlist(tlist(tint))),
+        #      [((l1,l2),list(map(lambda x,y: [x]+y,l1,l2)))
+        #       for _ in range(10)
+        #       for l1 in [randomList()]
+        #       for l2 in [[ randomList() for _ in range(len(l1)) ]]]),
+    ]
+
+    # Learning to filter
+    filterBootstrap = [
+        # Task("remove empty lists",
+        #      arrow(tlist(tlist(tbool)), tlist(tlist(tbool))),
+        #      [((ls,), [l for l in ls if len(l) > 0])
+        #       for _ in range(10)
+        #       for ls in [[[flip() for _ in range(randint(0, 3))]
+        #                   for _ in range(4)]]])
+        # Task("remove non 0s",
+        #      arrow(tlist(tint), tlist(tint)),
+        #      [((xs,), filter(lambda x: x == 0, xs))
+        #       for _ in range(10)
+        #       for xs in [[ randint(0,3) for _ in range(5) ]] ]),
+        Task("remove 0s",
+             arrow(tlist(tint), tlist(tint)),
+             [((xs,), [x for x in xs if x != 0])
+              for _ in range(10)
+              for xs in [[randint(0, 3) for _ in range(5)]]]),
+        Task("remove non-positives",
+             arrow(tlist(tint), tlist(tint)),
+             [((xs,), [x for x in xs if not (x > 1)])
+              for _ in range(10)
+              for xs in [[randint(0, 3) for _ in range(5)]]]),
+    ]
+
+    return lengthBootstrap + filterBootstrap + \
+        unfoldBootstrap + arrayBootstrap + foldBootstrap + mapBootstrap + zipBootstrap
+
+
+def bonusListProblems():
+    # Taken from https://www.ijcai.org/Proceedings/75/Papers/037.pdf
+    # These problems might be a lot easier if we do not use numbers
+    def randomList(lb=None, ub=None):
+        if lb is None:
+            lb = 2
+        if ub is None:
+            ub = 5
+        return [randint(0, 5) for _ in range(randint(lb, ub))]
+
+    bonus = [
+        Task(
+            "pair reverse", arrow(tlist(tint), tlist(tint)),
+            [((x,), [x[j + (1 if j % 2 == 0 else -1)]
+                     for j in range(len(x))])
+             for _ in range(5)
+             for x in [randomList(10, 10)]]
+        ),
+        Task(
+            "duplicate each element", arrow(tlist(tint), tlist(tint)),
+            [((x,), [a for z in x for a in [z] * 2])
+             for _ in range(5)
+             for x in [randomList(4, 6)]]
+        ),
+        Task(
+            "reverse duplicate each element", arrow(tlist(tint), tlist(tint)),
+            [((x,), [a for z in reversed(x) for a in [z] * 2])]
+        ),
+    ]
+    return bonus
+
+def sortBootstrap():
+    # These tasks have as their goal the learning of (1) filter, and
+    # (2) sort, which uses filter.
+    def flip(): return random() > 0.5
+    def randomList(lb=None, ub=None):
+        if lb is None:
+            lb = 2
+        if ub is None:
+            ub = 5
+        return [randint(0, 10) for _ in range(randint(lb, ub))]
+    def randomBooleanList():
+        return [flip() for _ in range(randint(4, 7))]
+    def removeDuplicates(l):
+        if len(l) == 0: return l
+        return [l[0]] + removeDuplicates([ z for z in l if z != l[0] ])
+    
+    filterBootstrap = [
+        # Task("remove empty lists",
+        #      arrow(tlist(tlist(tbool)), tlist(tlist(tbool))),
+        #      [((ls,), [l for l in ls if len(l) > 0])
+        #       for _ in range(10)
+        #       for ls in [[[flip() for _ in range(randint(0, 3))]
+        #                   for _ in range(4)]]]),
+        # Task("remove non 0s",
+        #      arrow(tlist(tint), tlist(tint)),
+        #      [((xs,), filter(lambda x: x == 0, xs))
+        #       for _ in range(10)
+        #       for xs in [[ randint(0,3) for _ in range(5) ]] ]),
+        Task("remove 0s",
+             arrow(tlist(tint), tlist(tint)),
+             [((xs,), [x for x in xs if x != 0])
+              for _ in range(10)
+              for xs in [[randint(0, 3) for _ in range(5)]]]),
+        # Task("remove primes",
+        #      arrow(tlist(tint), tlist(tint)),
+        #      [((xs,), [x for x in xs if not (x in {2,3,5,7,11,13,17,19,23})])
+        #       for _ in range(10)
+        #       for xs in [[randint(0, 20) for _ in range(7)]]]),
+        Task("remove squares",
+             arrow(tlist(tint), tlist(tint)),
+             [((xs,), [x for x in xs if not (int(x**0.5)**2 == x)])
+              for _ in range(10)
+              for xs in [[randint(0, 20) for _ in range(7)]]]),
+        Task("remove > 1",
+             arrow(tlist(tint), tlist(tint)),
+             [((xs,), [x for x in xs if not (x > 1)])
+              for _ in range(10)
+              for xs in [[randint(0, 5) for _ in range(7)]]]),
+    ]
+
+    # Needed for selection sort
+    minimumBootstrap = [
+        Task("min2", arrow(tint,tint,tint),
+             [((x,y),min(x,y))
+              for x in range(4)
+              for y in range(4) ]),
+        Task("minimum of list", arrow(tlist(tint),tint),
+             [((l,),min(l))
+              for _ in range(15) 
+              for l in [randomList()] ])
+    ]
+
+    appendBootstrap = [
+        Task("append bool", arrow(tlist(tbool), tlist(tbool), tlist(tbool)),
+             [((x, y), x + y)
+              for _ in range(10)
+              for [x, y] in [[randomBooleanList(), randomBooleanList()]]]),
+        Task("append int", arrow(tlist(tint), tlist(tint), tlist(tint)),
+             [((x, y), x + y)
+              for _ in range(10)
+              for [x, y] in [[randomList(), randomList()]]])
+    ]
+
+    insertionBootstrap = [
+        Task("filter greater than or equal", arrow(tint,tlist(tint),tlist(tint)),
+             [((x,l), [y for y in l if y >= x ])
+              for _ in range(15) 
+              for x in [randint(0,5)]
+              for l in [randomList()] ]),
+        Task("filter less than", arrow(tint,tlist(tint),tlist(tint)),
+             [((x,l), [y for y in l if y < x ])
+              for _ in range(15) 
+              for x in [randint(0,5)]
+              for l in [randomList()] ]),
+        Task("insert into sorted list (I)", arrow(tint,tlist(tint),tlist(tint)),
+             [((x,l), [y for y in l if y < x ] + [x] + [y for y in l if y >= x ])
+              for _ in range(15) 
+              for x in [randint(0,5)]
+              for _l in [randomList()]
+              for l in [sorted(_l)] ]),
+        Task("insert into sorted list (II)", arrow(tint,tlist(tint),tlist(tint)),
+             [((x,l), [y for y in l if y < x ] + [x] + [y for y in l if y >= x ])
+              for _ in range(15) 
+              for x in [randint(0,5)]
+              for l in [randomList()] ])
+    ]
+
+
+    sortTask = [
+        Task("sort-and-deduplicate", arrow(tlist(tint),tlist(tint)),
+             [((l,),list(sorted(l)))
+              for _ in range(15)
+              for l in [removeDuplicates(randomList())]
+             ])]
+
+    slowSort = [
+        Task("+1 maximum list", arrow(tlist(tint), tint),
+             [((l,),max(l) + 1)
+              for _ in range(15)
+              for l in [randomList()] ]),
+        Task("range +1 maximum list", arrow(tlist(tint), tlist(tint)),
+             [((l,),list(range(max(l) + 1)))
+              for _ in range(15)
+              for l in [randomList()] ]),
+        ]
+        
+
+    tasks = sortTask + slowSort
+    for t in tasks: t.mustTrain = True
+    return tasks
+    
+
+def exportTasks():
+    import sys
+    import pickle as pickle
+
+    n_examples = 15
+    if len(sys.argv) > 1:
+        n_examples = int(sys.argv[1])
+
+    eprint("Downloading and generating dataset")
+    tasks = sorted(make_list_tasks(n_examples), key=lambda t: t.name)
+    eprint("Got {} list tasks".format(len(tasks)))
+
+    with open("data/list_tasks.pkl", "w") as f:
+        pickle.dump(tasks, f)
+    eprint("Wrote list tasks to data/list_tasks.pkl")
+
+
+if __name__ == "__main__":
+    import json
+    def retrieveJSONTasks(filename, features=False):
+        """
+        For JSON of the form:
+            {"name": str,
+             "type": {"input" : bool|int|list-of-bool|list-of-int,
+                      "output": bool|int|list-of-bool|list-of-int},
+             "examples": [{"i": data, "o": data}]}
+        """
+        with open(filename, "r") as f:
+            loaded = json.load(f)
+        TP = {
+            "bool": tbool,
+            "int": tint,
+            "list-of-bool": tlist(tbool),
+            "list-of-int": tlist(tint),
+        }
+        return [Task(
+            item["name"],
+            arrow(TP[item["type"]["input"]], TP[item["type"]["output"]]),
+            [((ex["i"],), ex["o"]) for ex in item["examples"]],
+            features=(None if not features else list_features(
+                [((ex["i"],), ex["o"]) for ex in item["examples"]])),
+            cache=False,
+        ) for item in loaded]
+    for t in retrieveJSONTasks("data/list_tasks.json") + sortBootstrap() + make_list_bootstrap_tasks():
+        print(t.describe())
+        print()
+    # exportTasks()

dreamcoder/domains/logo/logoPrimitives.py

@@ -0,0 +1,41 @@
+from dreamcoder.program import Primitive, Program
+from dreamcoder.type import arrow, baseType, tint
+
+turtle = baseType("turtle")
+tstate = baseType("tstate")
+tangle = baseType("tangle")
+tlength = baseType("tlength")
+
+primitives = [
+    Primitive("logo_UA", tangle, ""),
+    Primitive("logo_UL", tlength, ""),
+
+    Primitive("logo_ZA", tangle, ""),
+    Primitive("logo_ZL", tlength, ""),
+
+    Primitive("logo_DIVA",  arrow(tangle,tint,tangle), ""),
+    Primitive("logo_MULA",  arrow(tangle,tint,tangle), ""),
+    Primitive("logo_DIVL",  arrow(tlength,tint,tlength), ""),
+    Primitive("logo_MULL",  arrow(tlength,tint,tlength), ""),
+
+    Primitive("logo_ADDA",  arrow(tangle,tangle,tangle), ""),
+    Primitive("logo_SUBA",  arrow(tangle,tangle,tangle), ""),
+    # Primitive("logo_ADDL",  arrow(tlength,tlength,tlength), ""),
+    # Primitive("logo_SUBL",  arrow(tlength,tlength,tlength), ""),
+
+    # Primitive("logo_PU",  arrow(turtle,turtle), ""),
+    # Primitive("logo_PD",  arrow(turtle,turtle), ""),
+    Primitive("logo_PT", arrow(arrow(turtle,turtle),arrow(turtle,turtle)), None),
+    Primitive("logo_FWRT",  arrow(tlength,tangle,turtle,turtle), ""),
+    Primitive("logo_GETSET",  arrow(arrow(turtle,turtle),turtle,turtle), "")
+] + [
+    Primitive("logo_IFTY", tint, ""),
+    Primitive("logo_epsA", tangle, ""),
+    Primitive("logo_epsL", tlength, ""),
+    Primitive("logo_forLoop", arrow(tint, arrow(tint, turtle, turtle), turtle, turtle), "ERROR: python has no way of expressing this hence you shouldn't eval on this"),
+] + [Primitive(str(j), tint, j) for j in range(10)]
+
+if __name__ == "__main__":
+    expr_s = "(lambda (logo_forLoop 3 (lambda (lambda (logo_GET (lambda (logo_FWRT (logo_S2L (logo_I2S 1)) (logo_S2A (logo_I2S 0)) (logo_SET $0 (logo_FWRT (logo_S2L eps) (logo_DIVA (logo_S2A (logo_I2S 2)) (logo_I2S 3)) ($1)))))))) ($0)))"
+    x = Program.parse(expr_s)
+    print(x)

dreamcoder/domains/logo/main.py

@@ -0,0 +1,450 @@
+from collections import OrderedDict
+import datetime
+import json
+import os
+import pickle
+import random as random
+import subprocess
+import sys
+import time
+
+try:
+    import numpy as np
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+except:
+    print("WARNING: Could not import torch. This is only okay when doing pypy compression.",
+          file=sys.stderr)
+
+from dreamcoder.domains.logo.makeLogoTasks import makeTasks, montageTasks, drawLogo
+from dreamcoder.domains.logo.logoPrimitives import primitives, turtle, tangle, tlength
+from dreamcoder.dreamcoder import ecIterator
+from dreamcoder.grammar import Grammar
+from dreamcoder.program import Program
+try:
+    from dreamcoder.recognition import variable, maybe_cuda
+except:
+    print("WARNING: Could not import recognition. This is only okay when doing pypy compression.",
+          file=sys.stderr)
+from dreamcoder.task import Task
+from dreamcoder.type import arrow
+from dreamcoder.utilities import eprint, testTrainSplit, loadPickle
+
+
+def animateSolutions(allFrontiers):
+    programs = []
+    filenames = []
+    for n,(t,f) in enumerate(allFrontiers.items()):
+        if f.empty: continue
+
+        programs.append(f.bestPosterior.program)
+        filenames.append(f"/tmp/logo_animation_{n}")
+        
+    drawLogo(*programs, pretty=True, smoothPretty=True, resolution=128, animate=True,
+             filenames=filenames)
+        
+        
+    
+def dreamFromGrammar(g, directory, N=100):
+    if isinstance(g,Grammar):
+        programs = [ p
+                     for _ in range(N)
+                     for p in [g.sample(arrow(turtle,turtle),
+                                        maximumDepth=20)]
+                     if p is not None]
+    else:
+        programs = g
+    drawLogo(*programs,
+             pretty=False, smoothPretty=False,
+             resolution=512,
+             filenames=[f"{directory}/{n}.png" for n in range(len(programs)) ],
+             timeout=1)
+    drawLogo(*programs,
+             pretty=True, smoothPretty=False,
+             resolution=512,
+             filenames=[f"{directory}/{n}_pretty.png" for n in range(len(programs)) ],
+             timeout=1)
+    drawLogo(*programs,
+             pretty=False, smoothPretty=True,
+             resolution=512,
+             filenames=[f"{directory}/{n}_smooth_pretty.png" for n in range(len(programs)) ],
+             timeout=1)
+    for n,p in enumerate(programs):
+        with open(f"{directory}/{n}.dream","w") as handle:
+            handle.write(str(p))        
+    
+try:
+    class Flatten(nn.Module):
+        def __init__(self):
+            super(Flatten, self).__init__()
+
+        def forward(self, x):
+            return x.view(x.size(0), -1)
+
+
+    class LogoFeatureCNN(nn.Module):
+        special = "LOGO"
+
+        def __init__(self, tasks, testingTasks=[], cuda=False, H=64):
+            super(LogoFeatureCNN, self).__init__()
+
+            self.sub = prefix_dreams + str(int(time.time()))
+
+            self.recomputeTasks = False
+
+            def conv_block(in_channels, out_channels, p=True):
+                return nn.Sequential(
+                    nn.Conv2d(in_channels, out_channels, 3, padding=1),
+                    # nn.BatchNorm2d(out_channels),
+                    nn.ReLU(),
+                    # nn.Conv2d(out_channels, out_channels, 3, padding=1),
+                    # nn.ReLU(),
+                    nn.MaxPool2d(2))
+
+            self.inputImageDimension = 128
+            self.resizedDimension = 128
+            assert self.inputImageDimension % self.resizedDimension == 0
+
+            # channels for hidden
+            hid_dim = 64
+            z_dim = 64
+
+            self.encoder = nn.Sequential(
+                conv_block(1, hid_dim),
+                conv_block(hid_dim, hid_dim),
+                conv_block(hid_dim, hid_dim),
+                conv_block(hid_dim, hid_dim),
+                conv_block(hid_dim, hid_dim),
+                conv_block(hid_dim, z_dim),
+                Flatten()
+            )
+
+            self.outputDimensionality = 256
+
+
+
+
+        def forward(self, v):
+            assert len(v) == self.inputImageDimension*self.inputImageDimension
+            floatOnlyTask = list(map(float, v))
+            reshaped = [floatOnlyTask[i:i + self.inputImageDimension]
+                        for i in range(0, len(floatOnlyTask), self.inputImageDimension)]
+            v = variable(reshaped).float()
+            # insert channel and batch
+            v = torch.unsqueeze(v, 0)
+            v = torch.unsqueeze(v, 0)
+            v = maybe_cuda(v, next(self.parameters()).is_cuda)/256.
+            window = int(self.inputImageDimension/self.resizedDimension)
+            v = F.avg_pool2d(v, (window,window))
+            v = self.encoder(v)
+            return v.view(-1)
+
+        def featuresOfTask(self, t):  # Take a task and returns [features]
+            return self(t.highresolution)
+
+        def tasksOfPrograms(self, ps, types):
+            images = drawLogo(*ps, resolution=128)
+            if len(ps) == 1: images = [images]
+            tasks = []
+            for i in images:
+                if isinstance(i, str): tasks.append(None)
+                else:
+                    t = Task("Helm", arrow(turtle,turtle), [])
+                    t.highresolution = i
+                    tasks.append(t)
+            return tasks        
+
+        def taskOfProgram(self, p, t):
+            return self.tasksOfPrograms([p], None)[0]
+except:
+    pass
+
+def list_options(parser):
+    parser.add_argument("--proto",
+                        default=False,
+                        action="store_true",
+                        help="Should we use prototypical networks?")
+    parser.add_argument("--target", type=str,
+                        default=[],
+                        action='append',
+                        help="Which tasks should this try to solve")
+    parser.add_argument("--reduce", type=str,
+                        default=[],
+                        action='append',
+                        help="Which tasks should this try to solve")
+    parser.add_argument("--save", type=str,
+                        default=None,
+                        help="Filepath output the grammar if this is a child")
+    parser.add_argument("--prefix", type=str,
+                        default="experimentOutputs/",
+                        help="Filepath output the grammar if this is a child")
+    parser.add_argument("--dreamCheckpoint", type=str,
+                        default=None,
+                        help="File to load in order to get dreams")
+    parser.add_argument("--dreamDirectory", type=str,
+                        default=None,
+                        help="Directory in which to dream from --dreamCheckpoint")
+    parser.add_argument("--visualize",
+                        default=None, type=str)
+    parser.add_argument("--cost", default=False, action='store_true',
+                        help="Impose a smooth cost on using ink")
+    parser.add_argument("--split",
+                        default=1., type=float)
+    parser.add_argument("--animate",
+                        default=None, type=str)
+
+
+
+def outputDreams(checkpoint, directory):
+    from dreamcoder.utilities import loadPickle
+    result = loadPickle(checkpoint)
+    eprint(" [+] Loaded checkpoint",checkpoint)
+    g = result.grammars[-1]
+    if directory is None:
+        randomStr = ''.join(random.choice('0123456789') for _ in range(10))
+        directory = "/tmp/" + randomStr
+    eprint(" Dreaming into",directory)
+    os.system("mkdir  -p %s"%directory)
+    dreamFromGrammar(g, directory)
+
+def enumerateDreams(checkpoint, directory):
+    from dreamcoder.dreaming import backgroundHelmholtzEnumeration
+    from dreamcoder.utilities import loadPickle
+    result = loadPickle(checkpoint)
+    eprint(" [+] Loaded checkpoint",checkpoint)
+    g = result.grammars[-1]
+    if directory is None: assert False, "please specify a directory"
+    eprint(" Dreaming into",directory)
+    os.system("mkdir  -p %s"%directory)
+    frontiers = backgroundHelmholtzEnumeration(makeTasks(None,None), g, 100,
+                                               evaluationTimeout=0.01,
+                                               special=LogoFeatureCNN.special)()
+    print(f"{len(frontiers)} total frontiers.")
+    MDL = 0
+    def L(f):
+        return -list(f.entries)[0].logPrior
+    frontiers.sort(key=lambda f: -L(f))
+    while len(frontiers) > 0:
+        # get frontiers whose MDL is between [MDL,MDL + 1)
+        fs = []
+        while len(frontiers) > 0 and L(frontiers[-1]) < MDL + 1:
+            fs.append(frontiers.pop(len(frontiers) - 1))
+        if fs:
+            random.shuffle(fs)
+            print(f"{len(fs)} programs with MDL between [{MDL}, {MDL + 1})")
+
+            fs = fs[:500]
+            os.system(f"mkdir {directory}/{MDL}")
+            dreamFromGrammar([list(f.entries)[0].program for f in fs],
+                             f"{directory}/{MDL}")
+        MDL += 1
+
+def visualizePrimitives(primitives, export='/tmp/logo_primitives.png'):
+    from itertools import product
+    from dreamcoder.program import Index,Abstraction,Application
+    from dreamcoder.utilities import montageMatrix,makeNiceArray
+    from dreamcoder.type import tint
+    import scipy.misc
+    from dreamcoder.domains.logo.makeLogoTasks import parseLogo
+
+    angles = [Program.parse(a)
+              for a in ["logo_ZA",
+                        "logo_epsA",
+                        "(logo_MULA logo_epsA 2)",
+                        "(logo_DIVA logo_UA 4)",
+                        "(logo_DIVA logo_UA 5)",
+                        "(logo_DIVA logo_UA 7)",
+                        "(logo_DIVA logo_UA 9)",
+                        ] ]
+    specialAngles = {"#(lambda (lambda (logo_forLoop logo_IFTY (lambda (lambda (logo_FWRT (logo_MULL logo_UL 3) (logo_MULA $2 4) $0))) $1)))":
+                     [Program.parse("(logo_MULA logo_epsA 4)")]+[Program.parse("(logo_DIVA logo_UA %d)"%n) for n in [7,9] ]}
+    numbers = [Program.parse(n)
+               for n in ["1","2","5","7","logo_IFTY"] ]
+    specialNumbers = {"#(lambda (#(lambda (lambda (lambda (lambda (logo_forLoop $2 (lambda (lambda (logo_FWRT $5 (logo_DIVA logo_UA $3) $0))) $0))))) (logo_MULL logo_UL $0) 4 4))":
+                      [Program.parse(str(n)) for n in [1,2,3] ]}
+    distances = [Program.parse(l)
+                 for l in ["logo_ZL",
+                           "logo_epsL",
+                           "(logo_MULL logo_epsL 2)",
+                           "(logo_DIVL logo_UL 2)",
+                           "logo_UL"] ]
+    subprograms = [parseLogo(sp)
+                   for sp in ["(move 1d 0a)",
+                              "(loop i infinity (move (*l epsilonLength 4) (*a epsilonAngle 2)))",
+                              "(loop i infinity (move (*l epsilonLength 5) (/a epsilonAngle 2)))",
+                              "(loop i 4 (move 1d (/a 1a 4)))"]]
+
+    entireArguments = {"#(lambda (lambda (#(#(lambda (lambda (lambda (logo_forLoop $2 (lambda (lambda (logo_FWRT $2 $3 $0))))))) logo_IFTY) (logo_MULA (#(logo_DIVA logo_UA) $1) $0) (#(logo_MULL logo_UL) 3))))":
+                       [[Program.parse(str(x)) for x in xs ]
+                        for xs in [("3", "1", "$0"),
+                                   ("4", "1", "$0"),
+                                   ("5", "1", "$0"),
+                                   ("5", "3", "$0"),
+                                   ("7", "3", "$0")]]}
+    specialDistances = {"#(lambda (lambda (logo_forLoop 7 (lambda (lambda (#(lambda (lambda (lambda (#(lambda (lambda (lambda (logo_forLoop $2 (lambda (lambda (logo_FWRT $2 $3 $0))))))) 7 $1 $2 $0)))) $3 logo_epsA $0))) $0)))":
+                        [Program.parse("(logo_MULL logo_epsL %d)"%n) for n in range(5)]}
+    
+    matrix = []
+    for p in primitives:
+        if not p.isInvented: continue
+        t = p.tp
+        eprint(p,":",p.tp)
+        if t.returns() != turtle:
+            eprint("\t(does not return a turtle)")
+            continue
+
+        def argumentChoices(t):
+            if t == turtle:
+                return [Index(0)]
+            elif t == arrow(turtle,turtle):
+                return subprograms
+            elif t == tint:
+                return specialNumbers.get(str(p),numbers)
+            elif t == tangle:
+                return specialAngles.get(str(p),angles)
+            elif t == tlength:
+                return specialDistances.get(str(p),distances)
+            else: return []
+
+        ts = []
+        for arguments in entireArguments.get(str(p),product(*[argumentChoices(t) for t in t.functionArguments() ])):
+            eprint(arguments)
+            pp = p
+            for a in arguments: pp = Application(pp,a)
+            pp = Abstraction(pp)
+            i = np.reshape(np.array(drawLogo(pp, resolution=128)), (128,128))
+            if i is not None:
+                ts.append(i)
+            
+
+        if ts == []: continue
+
+        matrix.append(ts)
+        if len(ts) < 6: ts = [ts]
+        else: ts = makeNiceArray(ts)
+        r = montageMatrix(ts)
+        fn = "/tmp/logo_primitive_%d.png"%len(matrix)
+        eprint("\tExported to",fn)
+        scipy.misc.imsave(fn, r)
+        
+    matrix = montageMatrix(matrix)
+    scipy.misc.imsave(export, matrix)
+
+
+def main(args):
+    """
+    Takes the return value of the `commandlineArguments()` function as input and
+    trains/tests the model on LOGO tasks.
+    """
+
+    # The below legacy global statement is required since prefix_dreams is used by LogoFeatureCNN.
+    # TODO(lcary): use argument passing instead of global variables.
+    global prefix_dreams
+
+    # The below global statement is required since primitives is modified within main().
+    # TODO(lcary): use a function call to retrieve and declare primitives instead.
+    global primitives
+
+    visualizeCheckpoint = args.pop("visualize")
+    if visualizeCheckpoint is not None:
+        with open(visualizeCheckpoint,'rb') as handle:
+            primitives = pickle.load(handle).grammars[-1].primitives
+        visualizePrimitives(primitives)
+        sys.exit(0)
+
+    dreamCheckpoint = args.pop("dreamCheckpoint")
+    dreamDirectory = args.pop("dreamDirectory")
+
+    proto = args.pop("proto")
+
+    if dreamCheckpoint is not None:
+        #outputDreams(dreamCheckpoint, dreamDirectory)
+        enumerateDreams(dreamCheckpoint, dreamDirectory)
+        sys.exit(0)
+
+    animateCheckpoint = args.pop("animate")
+    if animateCheckpoint is not None:
+        animateSolutions(loadPickle(animateCheckpoint).allFrontiers)
+        sys.exit(0)
+        
+    target = args.pop("target")
+    red = args.pop("reduce")
+    save = args.pop("save")
+    prefix = args.pop("prefix")
+    prefix_dreams = prefix + "/dreams/" + ('_'.join(target)) + "/"
+    prefix_pickles = prefix + "/logo." + ('.'.join(target))
+    if not os.path.exists(prefix_dreams):
+        os.makedirs(prefix_dreams)
+    tasks = makeTasks(target, proto)
+    eprint("Generated", len(tasks), "tasks")
+
+    costMatters = args.pop("cost")
+    for t in tasks:
+        t.specialTask[1]["costMatters"] = costMatters
+        # disgusting hack - include whether cost matters in the dummy input
+        if costMatters: t.examples = [(([1]), t.examples[0][1])]
+
+    os.chdir("prototypical-networks")
+    subprocess.Popen(["python","./protonet_server.py"])
+    time.sleep(3)
+    os.chdir("..")
+
+
+    test, train = testTrainSplit(tasks, args.pop("split"))
+    eprint("Split tasks into %d/%d test/train" % (len(test), len(train)))
+    try:
+        if test: montageTasks(test,"test_")    
+        montageTasks(train,"train_")
+    except:
+        eprint("WARNING: couldn't generate montage. Do you have an old version of scipy?")
+
+    if red is not []:
+        for reducing in red:
+            try:
+                with open(reducing, 'r') as f:
+                    prods = json.load(f)
+                    for e in prods:
+                        e = Program.parse(e)
+                        if e.isInvented:
+                            primitives.append(e)
+            except EOFError:
+                eprint("Couldn't grab frontier from " + reducing)
+            except IOError:
+                eprint("Couldn't grab frontier from " + reducing)
+            except json.decoder.JSONDecodeError:
+                eprint("Couldn't grab frontier from " + reducing)
+
+    primitives = list(OrderedDict((x, True) for x in primitives).keys())
+    baseGrammar = Grammar.uniform(primitives, continuationType=turtle)
+
+    eprint(baseGrammar)
+
+    timestamp = datetime.datetime.now().isoformat()
+    outputDirectory = "experimentOutputs/logo/%s"%timestamp
+    os.system("mkdir -p %s"%outputDirectory)
+
+
+    generator = ecIterator(baseGrammar, train,
+                           testingTasks=test,
+                           outputPrefix="%s/logo"%outputDirectory,
+                           evaluationTimeout=0.01,
+                           **args)
+
+    r = None
+    for result in generator:
+        iteration = len(result.learningCurve)
+        dreamDirectory = "%s/dreams_%d"%(outputDirectory, iteration)
+        os.system("mkdir  -p %s"%dreamDirectory)
+        eprint("Dreaming into directory",dreamDirectory)
+        dreamFromGrammar(result.grammars[-1],
+                         dreamDirectory)
+        r = result
+
+    needsExport = [str(z)
+                   for _, _, z
+                   in r.grammars[-1].productions
+                   if z.isInvented]
+    if save is not None:
+        with open(save, 'w') as f:
+            json.dump(needsExport, f)

dreamcoder/domains/logo/makeLogoTasks.py

@@ -0,0 +1,777 @@
+# coding: utf8
+
+import os
+import random
+import sys
+
+from dreamcoder.domains.logo.logoPrimitives import primitives, turtle
+from dreamcoder.task import Task
+from dreamcoder.program import Abstraction, Application, Index, Program
+from dreamcoder.type import arrow
+from dreamcoder.utilities import eprint, jsonBinaryInvoke, random_seed, montage
+from dreamcoder.grammar import Grammar
+
+
+def drawLogo(*programs,
+             timeout=None,
+             resolution=None,
+             pretty=False, smoothPretty=False,
+             filenames=[],
+             animate=False,
+             cost=False):
+    message = {}
+    if pretty: message["pretty"] = pretty
+    if smoothPretty: message["smoothPretty"] = smoothPretty
+    if timeout: message["timeout"] = timeout
+    assert resolution is not None, "resolution not provided in drawLogo"
+    if isinstance(resolution, list):
+        assert len(resolution) == len(programs), "must provide a resolution for each program"
+    elif isinstance(resolution, int):
+        resolution = [resolution]*len(programs)
+    else: assert False
+    jobs = []
+    for p, size in zip(programs, resolution):
+        entry = {"program": str(p),
+                 "size": size}
+        if animate: entry["animate"] = True
+        if len(filenames) > 0:
+            entry["export"] = filenames[0]
+            filenames = filenames[1:]
+        jobs.append(entry)        
+    message["jobs"] = jobs
+    response = jsonBinaryInvoke("./logoDrawString", message)
+    if cost:
+        # include the cost and return tuples of (pixels, cost)
+        response = [programResponse if isinstance(programResponse,str) else (programResponse["pixels"], programResponse["cost"])
+                    for programResponse in response ]
+    else:
+        response = [programResponse if isinstance(programResponse,str) else programResponse["pixels"]
+                    for programResponse in response ]
+    if len(programs) == 1:            
+        return response[0]
+    return response
+
+def makeTasks(subfolders, proto):
+    return manualLogoTasks()
+
+def parseLogo(s):
+        
+    _ua = Program.parse("logo_UA")
+    _ul = Program.parse("logo_UL")
+
+    _za = Program.parse("logo_ZA")
+    _zl = Program.parse("logo_ZL")
+
+    _da = Program.parse("logo_DIVA")
+    _ma = Program.parse("logo_MULA")
+    _dl = Program.parse("logo_DIVL")
+    _ml = Program.parse("logo_MULL")
+
+    _aa = Program.parse("logo_ADDA")
+    _sa = Program.parse("logo_SUBA")
+    _al = None#Program.parse("logo_ADDL")
+    _sl = None#Program.parse("logo_SUBL")
+
+    _pu = None#Program.parse("logo_PU")
+    _pd = None#Program.parse("logo_PD")
+    _p = Program.parse("logo_PT")
+    _move = Program.parse("logo_FWRT")
+    _embed = Program.parse("logo_GETSET")
+
+    _addition = Program.parse("+")
+    _infinity = Program.parse("logo_IFTY")
+    _ea = Program.parse("logo_epsA")
+    _el = Program.parse("logo_epsL")
+    _loop = Program.parse("logo_forLoop")
+
+    from sexpdata import loads, Symbol
+    s = loads(s)
+    def command(k, environment, continuation):
+        assert isinstance(k,list)
+        if k[0] == Symbol("move"):
+            return Application(Application(Application(_move,
+                                                       expression(k[1],environment)),
+                                           expression(k[2],environment)),
+                               continuation)
+        if k[0] == Symbol("for") or k[0] == Symbol("loop"):
+            v = k[1]
+            b = expression(k[2], environment)
+            newEnvironment = [None, v] + environment
+            body = block(k[3:], newEnvironment, Index(0))
+            return Application(Application(Application(_loop,b),
+                                           Abstraction(Abstraction(body))),
+                               continuation)
+        if k[0] == Symbol("embed"):
+            body = block(k[1:], [None] + environment, Index(0))
+            return Application(Application(_embed,Abstraction(body)),continuation)
+        if k[0] == Symbol("p"):
+            body = block(k[1:], [None] + environment, Index(0))
+            return Application(Application(_p,Abstraction(body)),continuation)
+
+        assert False
+    def expression(e, environment):
+        for n, v in enumerate(environment):
+            if e == v: return Index(n)
+
+        if isinstance(e,int): return Program.parse(str(e))
+
+        mapping = {"1a": _ua,
+                   "1d": _ul, "1l": _ul,
+                   "0a": _za,
+                   "0d": _zl, "0l": _zl,
+                   "/a": _da,
+                   "/l": _dl, "/d": _dl,
+                   "*a": _ma,
+                   "*l": _ml, "*d": _ml,
+                   "+a": _aa,
+                   "+d": _al, "+l": _al,
+                   "-a": _sa,
+                   "-d": _sl, "-l": _sl,
+                   "+": _addition,
+                   "infinity": _infinity,
+                   "epsilonAngle": _ea,
+                   "epsilonDistance": _el,
+                   "epsilonLength": _el}
+        if e == float('inf'): return _infinity
+        for name, value in mapping.items():
+            if e == Symbol(name): return value
+            
+        assert isinstance(e,list), "not a list %s"%e
+        for name, value in mapping.items():
+            if e[0] == Symbol(name):
+                f = value
+                for argument in e[1:]:
+                    f = Application(f, expression(argument, environment))
+                return f
+        assert False
+        
+    def block(b, environment, continuation):
+        if len(b) == 0: return continuation
+        return command(b[0], environment, block(b[1:], environment, continuation))
+
+    try: return Abstraction(command(s, [], Index(0)))
+    except: return Abstraction(block(s, [], Index(0)))
+
+
+def manualLogoTask(name, expression, proto=False, needToTrain=False,
+                   supervise=False, lambdaCalculus=False):
+    p = Program.parse(expression) if lambdaCalculus else parseLogo(expression)
+    from dreamcoder.domains.logo.logoPrimitives import primitives
+    from dreamcoder.grammar import Grammar
+    g = Grammar.uniform(primitives, continuationType=turtle)
+    gp = Grammar.uniform(primitives)
+    try:
+        l = g.logLikelihood(arrow(turtle,turtle),p)
+        lp = gp.logLikelihood(arrow(turtle,turtle),p)
+        assert l >= lp
+        eprint(name,-l,"nats")
+        
+    except: eprint("WARNING: could not calculate likelihood of manual logo",p)
+
+    attempts = 0
+    while True:
+        [output, highresolution] = drawLogo(p, p, resolution=[28,128], cost=True)
+        if output == "timeout" or highresolution == "timeout":
+            attempts += 1
+        else:
+            break
+    if attempts > 0:
+        eprint(f"WARNING: Took {attempts} attempts to render task {name} within timeout")
+
+    cost = output[1]
+    output = output[0]
+    assert highresolution[1] == cost
+    highresolution = highresolution[0]
+            
+    shape = list(map(int, output))
+    highresolution = list(map(float, highresolution))
+    t = Task(name, arrow(turtle,turtle),
+             [(([0]), shape)])
+    t.mustTrain = needToTrain
+    t.proto = proto
+    t.specialTask = ("LOGO", {"proto": proto})
+    t.specialTask[1]["cost"] = cost*1.05
+
+    t.highresolution = highresolution
+
+    if supervise:
+        t.supervisedSolution = p
+
+    return t
+
+def dSLDemo():
+    n = 0
+    demos = []
+    def T(source):
+        demos.append(manualLogoTask(str(len(demos)), source,
+                                    lambdaCalculus="lambda" in source))
+    # this looks like polygons - verify and include
+    T("(#(lambda (lambda (#(lambda (lambda (#(lambda (lambda (lambda (logo_forLoop $0 (lambda (lambda (logo_FWRT $4 $3 $0))))))) $1 $0 logo_IFTY))) $1 (logo_DIVA logo_UA $0)))) (logo_MULL logo_UL 4) 3)")
+    T("(#(lambda (lambda (#(lambda (lambda (#(lambda (lambda (lambda (logo_forLoop $0 (lambda (lambda (logo_FWRT $4 $3 $0))))))) $1 $0 logo_IFTY))) $1 (logo_DIVA logo_UA $0)))) (logo_MULL logo_UL 6) 4)")
+    T("(#(lambda (lambda (#(lambda (lambda (#(lambda (lambda (lambda (logo_forLoop $0 (lambda (lambda (logo_FWRT $4 $3 $0))))))) $1 $0 logo_IFTY))) $1 (logo_DIVA logo_UA $0)))) (logo_MULL logo_UL 5) 5)")
+    T("(#(lambda (lambda (#(lambda (lambda (#(lambda (lambda (lambda (logo_forLoop $0 (lambda (lambda (logo_FWRT $4 $3 $0))))))) $1 $0 logo_IFTY))) $1 (logo_DIVA logo_UA $0)))) (logo_MULL logo_UL 3) 6)")
+    T("(#(lambda (lambda (#(lambda (lambda (#(lambda (lambda (lambda (logo_forLoop $0 (lambda (lambda (logo_FWRT $4 $3 $0))))))) $1 $0 logo_IFTY))) $1 (logo_DIVA logo_UA $0)))) (logo_MULL logo_UL 2) 7)")
+
+    # Spirals!
+    for spiralSize in [1,2,3,4,5]:
+        T(f"((lambda (logo_forLoop logo_IFTY (lambda (lambda (logo_FWRT (logo_MULL logo_epsL $1) (logo_MULA logo_epsA $2) $0))))) {spiralSize})")
+    for spiralSize in [5,6,7,8,9]:
+        #T(f"(lambda (#(lambda (logo_forLoop $0 (lambda (lambda (#(lambda (logo_FWRT (logo_MULL logo_UL $0) (logo_DIVA logo_UA 4))) $1 $0))))) {spiralSize} $0))")
+        T("(loop i " + str(spiralSize) + " (move (*d 1l i) (/a 1a 4)))")# (#(lambda (logo_forLoop $0 (lambda (lambda (#(lambda (logo_FWRT (logo_MULL logo_UL $0) (logo_DIVA logo_UA 4))) $1 $0))))) {spiralSize} $0))")
+
+    # CIRCLES
+    #(lambda (#(lambda (logo_forLoop 6 (lambda (lambda (#(lambda (lambda (logo_forLoop logo_IFTY (lambda (lambda (logo_FWRT $2 $3 $0)))))) logo_epsA (logo_MULL logo_epsL $2) $0))))) 6 $0))
+    for circleSize in [1,3,5,7,9]:
+        T(f"(lambda (#(lambda (logo_forLoop 6 (lambda (lambda (#(lambda (lambda (logo_forLoop logo_IFTY (lambda (lambda (logo_FWRT $2 $3 $0)))))) logo_epsA (logo_MULL logo_epsL $2) $0))))) {circleSize} $0))")
+    
+    T("(loop i 3 (move (*d 1l 3) (/a 1a 4)))")
+    T("(loop i 5 (move (*d 1l 5) (/a 1a 5)))")
+    T("(loop i infinity (move (*d epsilonDistance 5) (/a epsilonAngle 3)))")
+    T("(loop i infinity (move (*d epsilonDistance 9) (/a epsilonAngle 2)))")
+    T("(loop i infinity (move (*d epsilonLength i) (*a epsilonAngle 3)))")
+    T("(loop i 9 (move (*d 1l i) (/a 1a 4)))")
+    T("(move 1d 0a)")
+    T("(loop i infinity (move (*d epsilonLength 6) epsilonAngle))")
+    T("(loop i infinity (move (*d epsilonLength 8) epsilonAngle))")
+    T("(loop k 2 (loop i infinity (move (*d epsilonLength 4) epsilonAngle)))")
+    T("(loop k 2 (loop i infinity (move (*d epsilonLength 8) epsilonAngle)))")
+    T("(loop s 4 (move (*d 1d 3) (/a 1a 4)))")
+    T("(loop s 4 (move (*d 1d 6) (/a 1a 4)))")
+    T("""
+          (loop j 5
+          (move 0d (/a 1a 5))
+          (embed (loop i infinity
+          (move (*d epsilonLength 6) epsilonAngle))
+          (loop i infinity
+          (move (*d epsilonLength 6) epsilonAngle))))""")
+    T("""
+              (loop j 5
+              (embed (loop s 4 (move (*d 1d 3) (/a 1a 4))))
+              (move 0d (/a 1a 5)))""")
+    return demos
+
+def rotationalSymmetryDemo():
+    demos = []
+    def T(source):
+        demos.append(manualLogoTask(str(len(demos)), source))
+            
+    body = {"dashed": "(p (move 1d 0a)) (move 1d 0a) (p (move 1d 0a)) (move 1d 0a)",
+            "lonely circle": "(p (move (*d 1d 2) 0a)) (loop k 2 (loop i infinity (move (*d epsilonLength 2) epsilonAngle)))",
+            "square dashed": "(p (move 1d 0a)) (loop s 4 (move 1d (/a 1a 4)))",
+            "square": "(loop s 4 (move (*d 1d 2) (/a 1a 4)))",
+            "semicircle": "(loop i infinity (move (*d epsilonLength 4) epsilonAngle))"}
+    for name in body:
+        for n in [3,4,5,6,7]:
+            T("""
+              (loop j %d
+              (embed %s)
+              (move 0d (/a 1a %d)))"""%(n,body[name],n))
+    return demos
+              
+
+def manualLogoTasks():
+    tasks = []
+    def T(name, source, needToTrain=False, supervise=False):
+        tasks.append(manualLogoTask(name, source, supervise=supervise,
+                                    needToTrain=needToTrain))
+    if False:
+        for d,a,s in [('1l','0a','(loop i infinity (move epsilonLength epsilonAngle))'),
+                      ('epsilonLength','0a','(loop i infinity (move epsilonLength epsilonAngle))'),
+                      ('(*d 1l 3)','0a','(move 1l 0a)'),
+                      ('epsilonLength','0a','(move (*d 1l 2) 0a)'),
+                      ('(*d epsilonLength 9)','0a','(move epsilonLength 0a)'),
+                      ('(/d 1l 2)','0a','(move 1l 0a)')]:
+            #            'epsilonLength']:
+            # for a in ['epsilonAngle','0a']:
+            #     for s in ['(move 1l 0a)',
+            #               '(move epsilonLength 0a)',
+            #               '(loop i infinity (move epsilonLength epsilonAngle))']:
+            #         if d == 'epsilonLength' and s == '(move epsilonLength 0a)': continue
+            T("pu: %s/%s/%s"%(d,a,s),
+              """
+              (pu (move %s %s) pd %s)
+              """%(d,a,s))
+        return tasks
+
+    def slant(n):
+        return f"(move 0d (/a 1a {n}))"
+
+    for n,l,s in [(3,"1l",8),
+                  (4,"(*d 1d 3)",None),
+                  (5,"1l",None),
+                  (6,"(*d 1d 2)",5),
+                  (7,"1l",None),
+                  (8,"(/d 1d 2)",None)]:
+        T(f"{n}-gon {l}{'' if s is None else ' slanted '+str(s)}",
+          f"""
+          ({'' if s is None else slant(s)}
+           (loop i {n}
+            (move {l} (/a 1a {n}))))
+          """,
+          needToTrain=True)
+    for n,l,s in [(3,"(*d 1l 2)",None),
+                (4,"(*d 1d 4)",None),
+                (5,"(*d 1d 2)",None),
+                (6,"1l",None),
+                (7,"(*d 1d 3)",None),
+                (8,"1l",3)]:
+        T(f"{n}-gon {l}{'' if s is None else ' slanted '+str(s)}",
+          f"""
+          ({'' if s is None else slant(s)}
+           (loop i {n}
+            (move {l} (/a 1a {n}))))
+          """,
+          needToTrain=False)
+
+        
+
+    T("upwards", "((move 0d (/a 1a 4)) (move 1d 0a))",
+      needToTrain=True)
+    T("right angle", "((move (*d 1d 2) (/a 1a 4)) (move 1d 0a))",
+      needToTrain=True)
+    T("right angle epsilon", "((move epsilonLength (/a 1a 4)) (move epsilonLength 0a))",
+      needToTrain=True)
+
+    T("line segment", "(move 1d 0a)",
+      needToTrain=True)
+
+    T("square slanted by 2pi/3",
+      """((move 0d (/a 1a 3))
+      (loop k 4 (move 1d (/a 1a 4))))""",
+      needToTrain=True)
+    T("semicircle slanted by 2pi/5",
+      """((move 0d (/a 1a 5))
+      (loop i infinity
+      (move (*d epsilonLength 4) epsilonAngle)))""",
+      needToTrain=True)
+    T("Greek spiral slanted by 2pi/6",
+      """((move 0d (/a 1a 6))
+      (loop i 7 (move (*l 1l i) (/a 1a 4))))""",
+      needToTrain=True)
+    T("Hook slanted by 2pi/7",
+      """((move 0d (/a 1a 7))
+      (move 1d 0a)
+      (loop i infinity
+      (move (*d epsilonLength 4) epsilonAngle)))""",
+      needToTrain=True)
+    T("""slanted line""",
+      """((move 0d (/a 1a 8))
+      (move (*d 1l 3) 0a))""",
+      needToTrain=True)
+    
+
+    for i in [6,7,8,9]:
+        T("Greek spiral %d"%i,
+          """
+          (loop i %d
+          (move (*l 1l i) (/a 1a 4)))
+          """%i,
+          needToTrain=i in [7,8])
+    for i in [2,3,4,5]:
+        T("smooth spiral %d"%i,
+          """
+          (loop i infinity 
+          (move (*d epsilonLength i) (*a epsilonAngle %d)))
+          """%i,
+          needToTrain=i in [3,5])
+
+    T("smooth spiral 4 slanted by 2pi/2",
+      """
+          ((move 0d (/a 1a 2))
+      (loop i infinity 
+          (move (*d epsilonLength i) (*a epsilonAngle 4))))
+      """,
+      needToTrain=True)
+
+    for i in [3,5,7,9]:
+        T("star %d"%i,
+          """
+          (loop i %d (move (*d 1d 4) (-a (/a 1a 2) (/a (/a 1a 2) %s))))
+          """%(i,i),
+          needToTrain=i in [5,9])
+
+    T("leaf iteration 1.1",
+      """
+      (loop i infinity (move epsilonDistance (/a epsilonAngle 2)))
+      """,
+      needToTrain=True)
+    T("leaf iteration 1.2",
+      """
+      ((move 0d (/a 1a 2))
+      (loop i infinity (move epsilonDistance (/a epsilonAngle 2))))
+      """,
+      needToTrain=True)
+    T("leaf iteration 2.1",
+      """
+      (loop n 2
+      (loop i infinity (move epsilonDistance (/a epsilonAngle 2)))
+      (move 0d (/a 1a 4)))
+      """,
+      needToTrain=True)
+    T("leaf iteration 2.2",
+      """
+      ((move 0d (/a 1a 2))
+      (loop n 2
+      (loop i infinity (move epsilonDistance (/a epsilonAngle 2)))
+      (move 0d (/a 1a 4))))
+      """,
+      needToTrain=True)
+    for n in range(3,8):
+        T("flower %d"%n,
+          """
+          (loop j %d
+          (loop n 2
+          (loop i infinity (move epsilonDistance (/a epsilonAngle 2)))
+          (move 0d (/a 1a 4)))
+          (move 0d (/a 1a %d)))
+          """%(n,n),
+          needToTrain=n in range(3,5))        
+
+    for n in [5,6]:
+        T("staircase %d"%n,
+          """
+          (loop i %d
+          (move 1d (/a 1a 4))
+          (move 1d (/a 1a 4))
+          (move 0d (/a 1a 2)))
+          """%n,
+          needToTrain=n in [5])
+
+    for n in range(1,6):
+        T("blocks zigzag %d"%n,
+          """
+          (loop i %d
+          (move 1d (/a 1a 4)) (move 1d (/a 1a 4))
+          (move 1d (+a (/a 1a 2) (/a 1a 4))) (move 1d (+a (/a 1a 2) (/a 1a 4))))
+          """%n,
+          needToTrain=n in [1,2,3])
+    for n in [3,4]:#range(1,5):
+        T("diagonal zigzag %d"%n,
+          """
+          ((move 0d (/a 1a 8))
+          (loop i %d
+          (move 1d (/a 1a 4)) 
+          (move 1d (+a (/a 1a 2) (/a 1a 4)))))
+          """%n,
+          needToTrain=n == 4)
+
+    
+
+    for n in [1,2,3,4,5,6]:
+        T("right semicircle of size %d"%n,
+          """
+          (loop i infinity
+          (move (*d epsilonLength %d) (-a 0a epsilonAngle)))
+          """%n,
+          needToTrain=n%2 == 0)
+        T("left semicircle of size %d"%n,
+          f"""
+          ({'' if n != 1 else slant(8)}
+           (loop i infinity
+            (move (*d epsilonLength {n}) epsilonAngle)))
+          """,
+          needToTrain=n%2 == 1)
+        T("circle of size %d"%n,
+              """
+              ((loop i infinity
+              (move (*d epsilonLength %d) epsilonAngle))
+              (loop i infinity
+              (move (*d epsilonLength %d) epsilonAngle)))
+              """%(n,n),
+          needToTrain=n in [1,4,3,5,6])
+
+    for n in [5,6]:
+        T("%d enclosed circles"%n,
+          """
+          (loop j %d
+          (loop i infinity
+          (move (*d epsilonLength j) epsilonAngle))
+          (loop i infinity
+          (move (*d epsilonLength j) epsilonAngle)))"""%n,
+          needToTrain=n == 5)
+
+    for n,l in [(4,2),
+                (5,3),
+                (6,4),
+                (3,1)]:
+        T("%d-circle flower l=%d"%(n,l),
+          """
+          (loop j %d
+          (move 0d (/a 1a %d))
+          (embed (loop i infinity
+          (move (*d epsilonLength %d) epsilonAngle))
+          (loop i infinity
+          (move (*d epsilonLength %d) epsilonAngle))))"""%(n,n,l,l),
+          needToTrain=(n,l) in [(6,4),(3,1)])
+
+    for n,l in [(3,1),(2,2),(1,3),
+                (2,1),(1,2),(1,1)]:
+        T("%d-semicircle sequence L=%d"%(n,l),
+          """
+          (loop j %d
+          (loop i infinity
+          (move (*d epsilonLength %d) epsilonAngle))
+          (loop i infinity
+          (move (*d epsilonLength %d) (-a 0a epsilonAngle))))
+          """%(n,l,l),
+          needToTrain=(n,l) in [(3,1),(2,2),(1,3)])
+
+    for n,l in [(2,"1d"),
+                (3,"1d")]:
+        T("row of %d circles"%n,
+          """
+          (loop j %d
+          (embed (loop k 2 (loop i infinity (move epsilonLength epsilonAngle))))
+          (p (move %s 0a)))"""%(n,l),
+          needToTrain=n == 2)
+    for n,l in [(2,"1d"),
+                (3,"1d")]:
+        T("row of %d lines"%n,
+          """
+          (loop j %d
+          (move 1d 0a)
+          (p (move %s 0a)))"""%(n,l),
+          needToTrain=n == 2)
+    T("line next to semicircle",
+      """
+      ((move 1d 0a) (p (move 1d 0a)) (loop i infinity (move epsilonLength epsilonAngle)))
+      """,
+      needToTrain=True)
+    for n,l in [(3,"(/d 1d 2)"),
+                (4,"(/d 1d 3)")]:
+        T("%d dashed lines of size %s"%(n,l),
+          """(loop i %d (p (move 1d 0a)) (move %s 0a))"""%(n,l),
+          needToTrain=n == 3)
+    T("broken circle",
+      """
+      ((loop i infinity (move epsilonLength epsilonAngle)) (p (move 1d 0a)) (loop i infinity (move epsilonLength epsilonAngle)))
+      """,
+      needToTrain=True)
+    T("circle next to semicircle",
+      """
+      ((loop i infinity (move epsilonLength epsilonAngle))
+      (loop i infinity (move epsilonLength epsilonAngle))
+      (p (move 1d 0a))
+      (loop i infinity (move epsilonLength epsilonAngle)))
+      """,
+      needToTrain=True)
+    T("semicircle next to square",
+      """
+      ((loop i infinity (move epsilonLength epsilonAngle))
+      (p (move 1d 0a))
+      (loop i infinity (move 1d (/a 1a 4))))
+      """,
+      needToTrain=False)
+    T("circle next to square",
+      """
+      ((loop i infinity (move epsilonLength epsilonAngle))
+      (loop i infinity (move epsilonLength epsilonAngle))
+      (p (move 1d 0a))
+      (loop i infinity (move 1d (/a 1a 4))))
+      """,
+      needToTrain=False)
+    T("circle next to line",
+      """
+      ((loop i infinity (move epsilonLength epsilonAngle))
+      (loop i infinity (move epsilonLength epsilonAngle))
+      (p (move 1d 0a))
+      (move 1d 0a))
+      """,
+      needToTrain=True)
+    T("line next to circle",
+      """
+      ((move 1d 0a)
+      (p (move 1d 0a))
+      (loop i infinity (move epsilonLength epsilonAngle))
+      (loop i infinity (move epsilonLength epsilonAngle))      
+      (move 1d 0a))
+      """,
+      needToTrain=True)
+    for n,l in [(4,"1d"),
+                (5,"1d")]:
+        T("row of %d dashes"%n,
+          """
+          (loop j %d
+          (embed (move 0d (/a 1a 4)) (move 1d 0a))
+          (p (move %s 0a)))"""%(n,l),
+          needToTrain=n == 4)        
+    for n,l in [(5,"1d"),(6,"1d")]:
+        T("row of %d semicircles"%n,
+          """
+          (loop j %d
+          (embed (loop i infinity (move epsilonLength epsilonAngle)))
+          (p (move %s 0a)))"""%(n,l),
+          needToTrain=n == 5)
+
+    with random_seed(42): # carefully selected for maximum entropy
+        for n in [3,4,5,6,7]:
+            body = {"empty": "(move 1d 0a)",
+                    "spiral": "(loop i infinity (move (*d epsilonLength i) (*a epsilonAngle 2)))",
+                    "dashed": "(p (move 1d 0a)) (move 1d 0a)",
+                    "circle": "(move 1d 0a) (loop k 2 (loop i infinity (move epsilonLength epsilonAngle)))",
+                    "lonely circle": "(p (move 1d 0a)) (loop k 2 (loop i infinity (move epsilonLength epsilonAngle)))",
+                    "square dashed": "(p (move 1d 0a)) (loop s 4 (move 1d (/a 1a 4)))",
+                    "square": "(move 1d 0a) (loop s 4 (move 1d (/a 1a 4)))",
+                    "close large semicircle": "(loop i infinity (move (*d epsilonLength 2) epsilonAngle))",
+                    "close semicircle": "(loop i infinity (move epsilonLength epsilonAngle))",
+                    "semicircle": "(move 1d 0a) (loop i infinity (move epsilonLength epsilonAngle))",
+                    "double dashed": "(p (move 1d 0a)) (move 1d 0a) (p (move 1d 0a)) (move 1d 0a)",
+                    "Greek": "(loop i 3 (move (*l 1l i) (/a 1a 4)))"}
+            for name in body:
+                if name == "spiral" and n not in [3,5]: continue
+                if name == "square" and n not in [5,3,6,7]: continue
+                if name == "semicircle" and n not in [5,3,4,6]: continue
+                if name == "Greek" and n not in [3,5]: continue
+                if name == "double dashed" and n not in [6,4,3]: continue
+                
+                mustTrain = False
+
+                mustTrain = mustTrain or (n == 3 and name == "Greek")
+                mustTrain = mustTrain or (n == 7 and name == "empty")
+                mustTrain = mustTrain or (n == 5 and name == "dashed")
+                mustTrain = mustTrain or (n == 7 and name == "circle")
+                mustTrain = mustTrain or (n == 6 and name == "circle")
+                mustTrain = mustTrain or (n == 6 and name == "lonely circle")
+                mustTrain = mustTrain or (n == 5 and name == "square")
+                mustTrain = mustTrain or (n == 7 and name == "square")
+                mustTrain = mustTrain or (n == 5 and name == "semicircle")
+                mustTrain = mustTrain or (n == 3 and name == "square dashed")
+                mustTrain = mustTrain or (n == 6 and name == "close semicircle")
+                mustTrain = mustTrain or (n == 5 and name == "close large semicircle")
+                mustTrain = mustTrain or (n == 3 and name == "spiral")
+                mustTrain = mustTrain or (n == 6 and name == "double dashed")
+                mustTrain = mustTrain or (n == 3 and name == "double dashed")
+                #mustTrain = mustTrain or (n == 6 and name == "empty")
+
+                #mustTrain = mustTrain or (random.random() < 0.07) # calibrated to give 70 training tasks
+                
+
+                # # cap number of super easy snowflakes
+                # if name == "empty" and n not in [7]: mustTrain = False
+                # if name == "dashed" and n not in [4]: mustTrain = False
+                
+
+                T("%d-%s snowflake"%(n,name),
+                  """
+                  (loop j %d
+                  (embed %s)
+                  (move 0d (/a 1a %d)))"""%(n,body[name],n),
+                  needToTrain=mustTrain)
+
+    for n in [3,4]:#2,3,4]:
+        T("%d-row of squares"%n,
+          """
+          (loop i %d
+          (embed (loop k 4 (move 1d (/a 1a 4))))
+          (move 1d 0a))
+          """%n,
+          needToTrain=n == 4)
+    T("2x2 grid",
+    """
+    (for x 2 (embed (for y 2
+       (embed (loop k 4 (move 1d (/a 1a 4))))
+       (move 1d 0a)))
+       (move 0d (/a 1a 4)) (move 1d (-a 0a (/a 1a 4))))
+    """)
+    T("slanted squares",
+      """
+      ((embed (loop k 4 (move 1d (/a 1a 4))))
+      (move 0d (/a 1a 8))
+      (loop k 4 (move 1d (/a 1a 4))))
+      """)
+    for l in range(1,6):
+        T("square of size %d"%l,
+          """
+          (for i 4
+          (move (*d 1d %d) (/a 1a 4)))
+          """%l,
+          needToTrain=l in range(4))
+    for n in [5,7]:
+        T("%d-concentric squares"%n,
+          """
+          (for i %d
+          (embed (loop j 4 (move (*d 1d i) (/a 1a 4)))))
+          """%n,
+          needToTrain=n == 5)
+    return tasks
+
+def montageTasks(tasks, prefix="", columns=None, testTrain=False):
+    import numpy as np
+    
+    w = 128
+    arrays = [t.highresolution for t in tasks]
+    for a in arrays:
+        assert len(a) == w*w
+
+    if testTrain:
+        arrays = [a for a,t in zip(arrays, tasks) if t.mustTrain ] + [a for a,t in zip(arrays, tasks) if not t.mustTrain ]
+        
+    arrays = [np.array([a[i:i + w]
+                        for i in range(0, len(a), w) ])
+              for a in arrays]
+    i = montage(arrays, columns=columns)
+
+    import scipy.misc        
+    scipy.misc.imsave('/tmp/%smontage.png'%prefix, i)
+    if testTrain:
+        trainingTasks = arrays[:sum(t.mustTrain for t in tasks)]
+        testingTasks = arrays[sum(t.mustTrain for t in tasks):]
+        random.shuffle(trainingTasks)
+        random.shuffle(testingTasks)
+        arrays = trainingTasks + testingTasks
+    else:
+        random.shuffle(arrays)
+    scipy.misc.imsave('/tmp/%srandomMontage.png'%prefix, montage(arrays, columns=columns))
+
+def demoLogoTasks():
+    import scipy.misc
+    import numpy as np
+
+    g0 = Grammar.uniform(primitives, continuationType=turtle)
+    eprint("dreaming into /tmp/dreams_0...")
+    N = 1000
+    programs = [ p
+                     for _ in range(N)
+                     for p in [g0.sample(arrow(turtle,turtle),
+                                         maximumDepth=20)]
+                     if p is not None]
+    os.system("mkdir  -p /tmp/dreams_0")
+    for n,p in enumerate(programs):
+        with open(f"/tmp/dreams_0/{n}.dream","w") as handle:
+            handle.write(str(p))
+    drawLogo(*programs, pretty=True, smoothPretty=False,
+             resolution=512,
+             filenames=[f"/tmp/dreams_0/{n}_pretty.png"
+                        for n in range(len(programs)) ],
+             timeout=1)
+    
+    if len(sys.argv) > 1:
+        tasks = makeTasks(sys.argv[1:],proto=False)
+    else:
+        tasks = makeTasks(['all'],proto=False)
+    montageTasks(tasks,columns=16,testTrain=True)
+    for n,t in enumerate(tasks):
+        a = t.highresolution
+        w = int(len(a)**0.5)
+        scipy.misc.imsave('/tmp/logo%d.png'%n, np.array([a[i:i+w]
+                                                         for i in range(0,len(a),w) ]))
+        logo_safe_name = t.name.replace("=","_").replace(' ','_').replace('/','_').replace("-","_") + ".png"
+        #os.system(f"convert /tmp/logo{n}.png -morphology Dilate Octagon /tmp/{logo_safe_name}")
+        os.system(f"convert /tmp/logo{n}.png -channel RGB -negate /tmp/{logo_safe_name}")
+    eprint(len(tasks),"tasks")
+    eprint(sum(t.mustTrain for t in tasks),"need to be trained on")
+
+    for t in dSLDemo():
+        a = t.highresolution
+        w = int(len(a)**0.5)
+        scipy.misc.imsave('/tmp/logoDemo%s.png'%t.name, np.array([a[i:i+w]
+                                                                  for i in range(0,len(a),w) ]))
+        os.system(f"convert /tmp/logoDemo{t.name}.png -morphology Dilate Octagon /tmp/logoDemo{t.name}_dilated.png")
+
+    tasks = [t for t in tasks if t.mustTrain ]
+    random.shuffle(tasks)
+    montageTasks(tasks[:16*3],"subset",columns=16)
+
+    montageTasks(rotationalSymmetryDemo(),"rotational")
+
+    

dreamcoder/domains/misc/RobustFillPrimitives.py

@@ -0,0 +1,308 @@
+#RobustFillPrimitives
+
+from dreamcoder.program import Primitive, prettyProgram
+from dreamcoder.grammar import Grammar
+from dreamcoder.type import tint, arrow, baseType #, t0, t1, t2
+
+from string import printable
+import re
+from collections import defaultdict
+
+#from functools import reduce
+
+
+disallowed = [
+    ("#", "hash"),
+    ("!", "bang"),
+    ("\"", "double_quote"),
+    ("$", "dollar"),
+    ("%", "percent"),
+    ("&", "ampersand"),
+    ("'", "single_quote"),
+    (")", "left_paren"),
+    ("(", "right_paren"),
+    ("*", "astrisk"),
+    ("+", "plus"),
+    (",", "comma"),
+    ("-", "dash"),
+    (".", "period"),
+    ("/", "slash"),
+    (":", "colon"),
+    (";", "semicolon"),
+    ("<", "less_than"),
+    ("=", "equal"),
+    (">", "greater_than"),
+    ("?", "question_mark"),
+    ("@", "at"),
+    ("[", "left_bracket"),
+    ("\\", "backslash"),
+    ("]", "right_bracket"),
+    ("^", "carrot"),
+    ("_", "underscore"),
+    ("`", "backtick"),
+    ("|", "bar"),
+    ("}", "right_brace"),
+    ("{", "left_brace"),
+    ("~", "tilde"),
+    (" ", "space"),
+    ("\t", "tab")
+]
+disallowed = dict(disallowed)
+delimiters = "&,.?!@()[]%{/}:;$#\"'"
+
+delim_dict = {disallowed[c]:c for c in delimiters}
+
+types = {}
+types["Number"] = r'\d+'
+types["Word"] = r'\w+'
+types["Alphanum"] = r'\w'
+types["PropCase"] = r'[A-Z][a-z]+'
+types["AllCaps"] = r'[A-Z]'
+types["Lower"] = r'[a-z]'
+types["Digit"] = r'\d'
+types["Char"] = r'.' 
+
+regexes = {name: re.escape(val) for name, val in delim_dict.items()}
+regexes = {**regexes, **types}
+
+tposition = baseType("position")
+tindex = baseType("index")
+tcharacter = baseType("character")
+tboundary = baseType("boundary")
+tregex = baseType("regex")
+tsubstr = baseType("substr")
+texpression = baseType("expression")
+tprogram = baseType("program")
+tnesting = baseType("nesting")
+ttype = baseType("type")
+tdelimiter = baseType("delimiter")
+
+def _substr(k1): return lambda k2: lambda string: string[k1:k2] #i think this is fine
+def _getspan(r1): 
+    return lambda i1: lambda b1: lambda r2: lambda i2: lambda b2: lambda string: \
+    string[
+    [m.end() for m in re.finditer(r1, string)][i1] if b1 == "End" else [m.start() for m in re.finditer(r1, string)][i1]:[m.end() for m in re.finditer(r2, string)][i2] if b2 == "End" else [m.start() for m in re.finditer(r2, string)][i2]
+    ]
+    #TODO format correctly
+def _getspan_const(r1): return lambda i1: lambda b1: lambda r2: lambda i2: lambda b2: (defaultdict(int, {r1:i1+1 if i1>=0 else abs(i1), r2:i2+1 if i2>=0 else abs(i2)}), max(i1+1 if i1>=0 else abs(i1), i2+1 if i2>=0 else abs(i2)))
+
+
+def _trim(string): 
+    assert False
+    return string
+
+def _replace(d1, d2): return lambda string: string.replace(d1,d2)
+
+def _getall(tp): return lambda string: ''.join(re.findall(tp, string))
+def _getfirst(tp, i): return lambda string: ''.join(re.findall(tp, string)[:i])
+def _gettoken(tp, i): return lambda string: re.findall(tp, string)[i]
+def _gettoken_const(tp, i): return defaultdict(int, {tp: i+1 if i>=0 else abs(i)}), i+1 if i>=0 else abs(i)
+
+def _getupto(reg): return lambda string: string[:[m.end() for m in re.finditer(reg, string)][0]]
+def _getfrom(reg): return lambda string: string[[m.end() for m in re.finditer(reg, string)][-1]:]
+
+def _concat2(expr1): return lambda expr2: lambda string: expr1(string) + expr2(string) #More concats plz
+def _concat1(expr): return lambda string: expr(string)
+def _concat_list(expr): return lambda program: lambda string: expr(string) + program(string)
+#i've decided that all of the things which are expressions should take tstring as last input and output a tstring. Thus, all requests are arrow(tstring, tstring) and we limit size with recursive depth
+"""
+todo: 
+- _trim
+- incorporate tcharacter 
+- constraints
+- format _getspan
+- figure out how to represent on top_level
+
+- flatten for nn
+- parse
+
+- robustfill_util
+- train dc model for robustfill
+- main_supervised_robustfill
+- evaluate_robustfill
+- sample_data
+
+
+- deal with escapes ... 
+
+constraints:
+elements, and number necessary, and lengths
+"""
+
+def robustFillPrimitives(max_len=100, max_index=5):
+    return [
+        #CPrimitive("concat2", arrow(texpression, texpression, tprogram), _concat2),
+        CPrimitive("concat1", arrow(texpression, tprogram), _concat1),
+        CPrimitive("concat_list", arrow(texpression, tprogram, tprogram), _concat_list),
+        #expressions
+        CPrimitive("Constant", arrow(tcharacter, texpression), lambda x: lambda y: x),  # add a constraint
+        CPrimitive("apply", arrow(tnesting, tsubstr, texpression), lambda n: lambda sub: lambda string: n(sub(string))),
+        CPrimitive("apply_n", arrow(tnesting, tnesting, texpression), lambda n1: lambda n2: lambda string: n1(n2(string))),
+        CPrimitive("expr_n", arrow(tnesting, texpression), lambda x: x),
+        CPrimitive("expr_f", arrow(tsubstr, texpression), lambda x: x)
+        ] + [
+        #substrings
+        CPrimitive("SubStr", arrow(tposition, tposition, tsubstr), _substr), # handled
+        CPrimitive("GetSpan", arrow(tregex, tindex, tboundary, tregex, tindex, tboundary, tsubstr), _getspan, _getspan_const)  #TODO constraint
+        ] + [
+        #nestings
+        CPrimitive("GetToken"+name+str(i), tnesting, _gettoken(tp,i), _gettoken_const(tp, i)) for name, tp in types.items() for i in range(-max_index, max_index)
+        ] + [
+        CPrimitive("ToCase_ProperCase", tnesting, lambda x: x.title(), (defaultdict(int, {r'[A-Z][a-z]+':1}), 1)),
+        CPrimitive("ToCase_AllCapsCase", tnesting, lambda x: x.upper(), (defaultdict(int, {r'[A-Z]':1}) ,1)),
+        CPrimitive("ToCase_LowerCase", tnesting, lambda x: x.lower(), (defaultdict(int, {r'[a-z]':1}), 1) )
+        ] + [
+        CPrimitive("Replace_"+name1+name2, tnesting, _replace(char1, char2), (defaultdict(int, {char1:1}), 1)) for name1, char1 in delim_dict.items() for name2, char2 in delim_dict.items() if char1 is not char2
+        ] + [
+        #CPrimitive("Trim", tnesting, _trim), #TODO
+        ] + [
+        CPrimitive("GetUpTo"+name, tnesting, _getupto(reg), (defaultdict(int, {reg:1} ),1)) for name, reg in regexes.items()
+        ] + [
+        CPrimitive("GetFrom"+name, tnesting, _getfrom(reg), (defaultdict(int, {reg:1} ),1)) for name, reg in regexes.items()
+        ] + [
+        CPrimitive("GetFirst_"+name+str(i), tnesting, _getfirst(tp, i), (defaultdict(int, {tp:i} ), i+1 if i>=0 else abs(i))) for name, tp in types.items() for i in list(range(-max_index,0))+ list(range(1,max_index+1))
+        ] + [ 
+        CPrimitive("GetAll_"+name, tnesting, _getall(reg),(defaultdict(int, {reg:1} ),1) ) for name, reg in types.items()
+        ] + [
+        #regexes
+        CPrimitive("type_to_regex", arrow(ttype, tregex), lambda x: x), #TODO also make disappear
+        CPrimitive("delimiter_to_regex", arrow(tdelimiter, tregex), lambda x: re.escape(x)) #TODO also make disappear
+        ] + [
+        #types
+        CPrimitive("Number", ttype, r'\d+', r'\d+'), #TODO
+        CPrimitive("Word", ttype, r'\w+', r'\w+'), #TODO
+        CPrimitive("Alphanum", ttype, r'\w', r'\w'), #TODO
+        CPrimitive("PropCase", ttype, r'[A-Z][a-z]+', r'[A-Z][a-z]+'), #TODO
+        CPrimitive("AllCaps", ttype, r'[A-Z]', r'[A-Z]'), #TODO
+        CPrimitive("Lower", ttype, r'[a-z]', r'[a-z]'), #TODO
+        CPrimitive("Digit", ttype, r'\d', r'\d'), #TODO
+        CPrimitive("Char", ttype, r'.', r'.') #TODO
+        ] + [
+        #Cases
+        # CPrimitive("ProperCase", tcase, .title()), #TODO
+        # CPrimitive("AllCapsCase", tcase, .upper()), #TODO
+        # CPrimitive("LowerCase", tcase, .lower()) #TODO
+        ] + [
+        #positions
+        CPrimitive("position"+str(i), tposition, i, (defaultdict(int), i+1 if i>=0 else abs(i)) ) for i in range(-max_len,max_len+1) #deal with indicies 
+        ] + [
+        #indices
+        CPrimitive("index"+str(i), tindex, i, i) for i in range(-max_index,max_index+1) #deal with indicies
+        ] + [
+        #characters
+        CPrimitive(i, tcharacter, i, (defaultdict(int, {i:1}),1) ) for i in printable[:-5] if i not in disallowed
+            ] + [
+        CPrimitive(name, tcharacter, char, (defaultdict(int, {char:1}), 1)) for char, name in disallowed.items() # NB: disallowed is reversed
+        ] + [
+        #delimiters
+        CPrimitive("delim_"+name, tdelimiter, char, char) for name, char in delim_dict.items()
+        ] + [
+        #boundaries
+        CPrimitive("End", tboundary, "End"),
+        CPrimitive("Start", tboundary, "Start")
+    ]
+
+
+
+def RobustFillProductions(max_len=100, max_index=5):
+    return [(0.0, prim) for prim in robustFillPrimitives(max_len=max_len, max_index=max_index)]
+
+
+def flatten_program(p):
+    string = p.show(False)
+    string = string.replace('(', '')
+    string = string.replace(')', '')
+    #remove '_fn' (optional)
+    string = string.split(' ')
+    string = list(filter(lambda x: x is not '', string))
+    return string
+
+
+
+
+def add_constraints(c1, c2=None):
+    if c2 is None:
+        return c1
+    d1, m1 = c1
+    d2, m2 = c2
+    min_size = max(m1, m2)
+    d = defaultdict(int)
+    for item in set(d1.keys()) | set(d2.keys()):
+        d[item] = max(d1[item], d2[item])
+    return d, min_size
+
+# class Constraint_prop:
+#     def application(self, p, environment):
+#         self.f.visit(self, environment)(self.x.visit(self, environment))
+#     def primitive(self, p, environment):
+#         return self.value
+
+class Constraint_prop:
+    def __init__(self):
+        pass
+
+    def application(self, p):
+        return p.f.visit(self)(p.x.visit(self))
+
+    def primitive(self, p):
+        return p.constraint
+
+    def execute(self, p):
+        return p.visit(self)
+    
+
+class CPrimitive(Primitive):
+    def __init__(self, name, ty, value, constraint=None):
+        #I have no idea why this works but it does ..... 
+        if constraint is None:
+            if len(ty.functionArguments())==0:
+                self.constraint = (defaultdict(int), 0)
+            elif len(ty.functionArguments())==1:
+                self.constraint = lambda x: x
+            elif len(ty.functionArguments())==2:
+                self.constraint = lambda x: lambda y: add_constraints(x,y)
+            else:
+                self.constraint = lambda x: x
+                for _ in range(len(ty.functionArguments()) - 1):
+                    self.constraint = lambda x: lambda y: add_constraints(x, self.constraint(y))
+        else: self.constraint = constraint
+        super(CPrimitive, self).__init__(name, ty, value)
+
+    #def __getinitargs__(self):
+    #    return (self.name, self.tp, self.value, None)
+
+    def __getstate__(self):
+        #print("self.name", self.name)
+        return self.name
+
+    def __setstate__(self, state):
+        #for backwards compatibility:
+        if type(state) == dict:
+            pass #do nothing, i don't need to load them if they are old...
+        else:
+            p = Primitive.GLOBALS[state]
+            self.__init__(p.name, p.tp, p.value, p.constraint) 
+
+
+
+if __name__=='__main__':
+    import time
+    CPrimitive("testCPrim", tint, lambda x: x, 17)
+    g = Grammar.fromProductions(RobustFillProductions())
+    print(len(g))
+    request = tprogram
+    p = g.sample(request)
+    print("request:", request)
+    print("program:")
+    print(prettyProgram(p))
+    s = 'abcdefg'
+    e = p.evaluate([])
+    #print("prog applied to", s)
+    #print(e(s))
+    print("flattened_program:")
+    flat = flatten_program(p)
+    print(flat)
+    t = time.time()    
+    constraints = Constraint_prop().execute(p)
+    print(time.time() - t)
+    print(constraints)

dreamcoder/domains/misc/algolispPrimitives.py

@@ -0,0 +1,508 @@
+#napsPrimitives.py
+from dreamcoder.program import Primitive, Program
+from dreamcoder.grammar import Grammar
+from dreamcoder.type import tlist, arrow, baseType #, t0, t1, t2
+
+#from functools import reduce
+
+
+#Internal TYPES:
+# NUMBER
+# BOOLEAN
+# NOTFUNCTYPE
+# Type
+# ANYTYPE
+
+#types
+tsymbol = baseType("symbol")
+#PROGRAM = SYMBOL = constant | argument | function_call | function | lambda
+tconstant = baseType("constant")
+tfunction = baseType("function")
+
+f = dict([("|||","triple_or"),
+("reduce","reduce"),
+("+","+"),
+("len","len"),
+("map","map"),
+("filter","filter"),
+("-","-"),
+("*","*"),
+("partial0","partial0"),
+("if","if"),
+("lambda1","lambda1"),
+("==","eq"),
+("range","range"),
+("digits","digits"),
+("slice","slice"),
+("reverse","reverse"),
+("lambda2","lambda2"),
+("deref","deref"),
+("partial1","partial1"),
+("/","div"),
+("<","less_than"),
+(">","greater_than"),
+("min","min"),
+("combine","combine"),
+("head","head"),
+("is_prime","is_prime"),
+("false","false"),
+("||","or"),
+("10","10"),
+("self","self"),
+("max","max"),
+("sort","sort"),
+("%","mod"),
+("invoke1","invoke1"),
+("!","bang"),
+("square","square"),
+("str_concat","str_concat"),
+("strlen","strlen"),
+("<=","leq"),
+("int-deref","int-deref"),
+("str_split","str_split"),
+("str_index","str_index"),
+("floor","floor"),
+("sqrt","sqrt"),
+("str_min","str_min"),
+("&&","AND"),
+("is_sorted","is_sorted"),
+("str_max","str_max"),
+(">=","geq")])
+
+fn_lookup = {
+**f
+}
+
+c = dict( 
+[("0","0"),
+("a","a"),
+("arg1","arg1"),
+("1","1"),
+("b","b"),
+("2","2"),
+("c","c"),
+("arg2","arg2"),
+("d","d"),
+("false","false"),
+("10","10"),
+("self","self"),
+("1000000000","1000000000"),
+("\"\"", "empty_str"),
+("e","e"),
+("40","40"),
+("f","f"),
+("\" \"", "space"),
+("g","g"),
+("\"z\"","z"),
+("true","true"),
+("h","h"),
+("i","i"),
+("j","j"),
+("k","k"),
+("l","l")]
+    )
+
+const_lookup = {
+        **c
+    }
+
+primitive_lookup = {**const_lookup, **fn_lookup}
+#Do i need arguments??
+
+def _fn_call(f):
+    #print("f", f)
+    def inner(sx):
+        #print("sx", sx)
+        if not type(sx) == list:
+            sx = [sx]
+        return [f] + sx
+    return lambda sx: inner(sx)
+
+def algolispPrimitives():
+    return [
+    Primitive("fn_call", arrow(tfunction, tlist(tsymbol), tsymbol), _fn_call),
+
+    Primitive("lambda1_call", arrow(tfunction, tlist(tsymbol), tsymbol), lambda f: lambda sx: ["lambda1", [f] + sx] if type(sx)==list else ["lambda1", [f] + [sx]] ),
+    Primitive("lambda2_call", arrow(tfunction, tlist(tsymbol), tsymbol), lambda f: lambda sx: ["lambda2", [f] + sx] if type(sx)==list else ["lambda2", [f] + [sx]] ),
+    #symbol converters:
+    # SYMBOL = constant | argument | function_call | function | lambda
+    Primitive("symbol_constant", arrow(tconstant, tsymbol), lambda x: x),
+    Primitive("symbol_function", arrow(tfunction, tsymbol), lambda x: x),
+    #list converters
+    Primitive('list_init_symbol', arrow(tsymbol, tlist(tsymbol)), lambda symbol: [symbol] ),
+    Primitive('list_add_symbol', arrow(tsymbol, tlist(tsymbol), tlist(tsymbol)), lambda symbol: lambda symbols: symbols + [symbol] if type(symbols) == list else [symbols] + [symbol])
+    ] + [
+    #functions:
+    Primitive(ec_name, tfunction, algo_name) for algo_name, ec_name in fn_lookup.items()
+    ] + [
+    #Constants
+    Primitive(ec_name, tconstant, algo_name) for algo_name, ec_name in const_lookup.items()
+    ]
+
+
+#for first pass, can just hard code vars and maps n stuff
+
+
+def algolispProductions():
+    return [(0.0, prim) for prim in algolispPrimitives()]
+
+
+algolisp_input_vocab = [
+"<S>",
+"</S>",
+"<UNK>",
+"|||",
+"(",
+")",
+"a",
+"b",
+"of",
+"the",
+"0",
+",",
+"arg1",
+"c",
+"and",
+"1",
+"reduce",
+"+",
+"int[]",
+"in",
+"given",
+"numbers",
+"int",
+"is",
+"len",
+"map",
+"digits",
+"d",
+"number",
+"array",
+"-",
+"filter",
+"to",
+"range",
+"are",
+"*",
+"partial0",
+"2",
+"if",
+"reverse",
+"that",
+"elements",
+"lambda1",
+"==",
+"an",
+"arg2",
+"values",
+"slice",
+"element",
+"lambda2",
+"deref",
+"you",
+"partial1",
+"e",
+"find",
+"your",
+"task",
+"compute",
+"among",
+"from",
+"consider",
+"first",
+"than",
+"value",
+"/",
+"what",
+"arrays",
+"with",
+"<",
+"length",
+">",
+"be",
+"min",
+"end",
+"sum",
+"one",
+"head",
+"f",
+"by",
+"combine",
+"segment",
+"coordinates",
+"not",
+"string",
+"is_prime",
+"false",
+"||",
+"at",
+"10",
+"half",
+"position",
+"self",
+"subsequence",
+"after",
+"such",
+"max",
+"prime",
+"sort",
+"let",
+"%",
+"longest",
+"inclusive",
+"which",
+"invoke1",
+"1000000000",
+"all",
+"positions",
+"!",
+"square",
+"its",
+"has",
+"reversed",
+"another",
+"less",
+"each",
+"\"\"",
+"order",
+"largest",
+"maximum",
+"g",
+"last",
+"smallest",
+"times",
+"strictly",
+"40",
+"smaller",
+"indexes",
+"str_concat",
+"strlen",
+"two",
+"starting",
+"<=",
+"on",
+"greater",
+"how",
+"many",
+"int-deref",
+"prefix",
+"bigger",
+"only",
+"str_split",
+"\" \"",
+"str_index",
+"can",
+"plus",
+"squared",
+"product",
+"strings",
+"floor",
+"sqrt",
+"before",
+"it",
+"concatenation",
+"index",
+"as",
+"define",
+"multiplied",
+"biggest",
+"rounded",
+"down",
+"string[]",
+"equal",
+"integer",
+"also",
+"based",
+"sorting",
+"replace",
+"becomes",
+"single",
+"digit",
+"characters",
+"keeping",
+"including",
+"h",
+"larger",
+"written",
+"divisible",
+"previous",
+"subarray",
+"mininum",
+"second",
+"middle",
+"same",
+"th",
+"median",
+"till",
+"integers",
+"sequence",
+"for",
+"indices",
+"between",
+"when",
+"doubled",
+"ending",
+"even",
+"multiply",
+"squares",
+"fibonacci",
+"exclusive",
+"odd",
+"keep",
+"whether",
+"minimum",
+"except",
+"letters",
+"appearing",
+"letter",
+"consecutive",
+"character",
+"factorial",
+"chosen",
+"start",
+"begin",
+"themselves",
+"\"z\"",
+"str_min",
+"remove",
+"present",
+"exist",
+"appear",
+"starts",
+"i",
+"located",
+"true",
+"&&",
+"found",
+"discarding",
+"is_sorted",
+"removing",
+"do",
+"increasing",
+"exceed",
+"ascending",
+"difference",
+"decremented",
+"existing",
+"alphabetically",
+"words",
+"added",
+"incremented",
+"backwards",
+"individual",
+"lexicographically",
+"separate",
+"abbreviation",
+"str_max",
+"increment",
+"consisting",
+"equals",
+"having",
+"discard",
+"descending",
+"decreasing",
+"sorted",
+"being",
+"where",
+"right",
+"there",
+"ordinal",
+"have",
+"s",
+"going",
+"'",
+"add",
+"space",
+"decrement",
+"those",
+"whitespaces",
+"spaces",
+"subtract",
+"remaining",
+"following",
+"or",
+"out",
+"ordered",
+"minimal",
+"itself",
+"symmetric",
+"read",
+"increases",
+"word",
+"immidiately",
+"excluding",
+"j",
+"omitting",
+"reads",
+"maximal",
+">=",
+"compare",
+"form",
+"absent",
+"missing",
+"cannot",
+"whose",
+"count",
+"lowest",
+"both",
+"ends",
+"beginning",
+"left",
+"mean",
+"average",
+"obtained",
+"writing",
+"result",
+"joining",
+"together",
+"increase",
+"highest",
+"comparing",
+"forms",
+"avg",
+"outside",
+"positive",
+"summed",
+"belonging",
+"lexicographical",
+"rest",
+"belong",
+"inclucing",
+"lexical",
+"alphabetical",
+"dictionary",
+"k",
+"negative",
+"lexicographic",
+"represents",
+"delete",
+"non",
+"l",
+"erase",
+"m",
+"comes",
+"up",
+"comparison",
+"during",
+"'s value is the largest inclusive, which is strictly less than maximum element in numbers from 1 to the element in `a` which'",
+"'s value is the biggest (inclusive), which is strictly less than maximum element of range from 1 to the element in `a` which'",
+"'s value is the highest, which is strictly less than maximum element among sequence of digits of the element in `a` which'"]
+
+
+if __name__ == "__main__":
+    #g = Grammar.uniform(deepcoderPrimitives())
+
+    g = Grammar.fromProductions(algolispProductions(), logVariable=.9)
+
+    #p=Program.parse("(lambda (fn_call filter (list_add_symbol (lambda1_call == (list_add_symbol 1 (list_init_symbol (fn_call mod ( list_add_symbol 2 (list_init_symbol arg1)) ))) ) (list_init_symbol $0)) )")
+    p=Program.parse("(lambda (fn_call filter (list_add_symbol (lambda1_call eq (list_add_symbol (symbol_constant 1) (list_init_symbol (fn_call mod ( list_add_symbol (symbol_constant 2) (list_init_symbol (symbol_constant arg1))) ))) ) (list_init_symbol (symbol_constant $0)))))")
+
+    print(p)
+
+    #tree = p.evaluate(["a"])
+    tree = p.evaluate([])
+    print(tree("a"))
+
+#
+
+  

dreamcoder/domains/misc/deepcoderPrimitives.py

@@ -0,0 +1,352 @@
+from dreamcoder.program import Primitive, prettyProgram
+from dreamcoder.grammar import Grammar
+from dreamcoder.type import tlist, tint, arrow, baseType #, t0, t1, t2
+
+#from functools import reduce
+
+
+#todo
+int_to_int = baseType("int_to_int")
+int_to_bool = baseType("int_to_bool")
+int_to_int_to_int = baseType("int_to_int_to_int")
+
+
+#deepcoderPrimitives
+Null = 300 #or perhaps something else, like "an integer outside the working range"?
+
+def _head(xs): return xs[0] if len(xs)>0 else Null
+def _tail(xs): return xs[-1] if len(xs)>0 else Null
+def _take(n): return lambda xs: xs[:n]
+def _drop(n): return lambda xs: xs[n:]
+def _access(n): return lambda xs: xs[n] if n>=0 and len(xs)>n else Null
+def _minimum(xs): return min(xs) if len(xs)>0 else Null
+def _maximum(xs): return max(xs) if len(xs)>0 else Null
+def _reverse(xs): return list(reversed(xs))
+def _sort(xs): return sorted(xs)
+def _sum(xs): return sum(xs)
+
+#higher order:
+def _map(f): return lambda l: list(map(f, l)) 
+def _filter(f): return lambda l: list(filter(f, l))
+def _count(f): return lambda l: len([x for x in l if f(x)])
+def _zipwith(f): return lambda xs: lambda ys: [f(x)(y) for (x, y) in zip(xs, ys)]
+def _scanl1(f): 
+    def _inner(xs):
+        ys = []
+        if len(xs) > 0:
+            ys.append(xs[0])
+            for i in range(1, len(xs)):
+                ys.append( f(ys[i-1])(xs[i]))
+        return ys
+    return _inner
+
+#int to int:
+def _succ(x): return x+1
+def _pred(x): return x-1
+def _double(x): return x*2
+def _half(x): return int(x/2)
+def _negate(x): return -x
+def _square(x): return x**2
+def _triple(x): return x*3
+def _third(x): return int(x/3)
+def _quad(x): return x*4
+def _quarter(x): return int(x/4)
+
+#int to bool:
+def _pos(x): return x>0
+def _neg(x): return x<0
+def _even(x): return x%2==0
+def _odd(x): return x%2==1
+
+#int to int to int:
+def _add(x): return lambda y: x + y
+def _sub(x): return lambda y: x - y
+def _mult(x): return lambda y: x * y
+def _min(x): return lambda y: _minimum([x,y])
+def _max(x): return lambda y: _maximum([x,y])
+
+def deepcoderPrimitives():
+    return [
+        Primitive("HEAD", arrow(tlist(tint), tint), _head), 
+        Primitive("LAST", arrow(tlist(tint), tint), _tail),
+        Primitive("TAKE", arrow(tint, tlist(tint), tlist(tint)), _take),
+        Primitive("DROP", arrow(tint, tlist(tint), tlist(tint)), _drop),
+        Primitive("ACCESS", arrow(tint, tlist(tint), tint), _access),
+        Primitive("MINIMUM", arrow(tlist(tint), tint), _minimum),
+        Primitive("MAXIMUM", arrow(tlist(tint), tint), _maximum),
+        Primitive("REVERSE", arrow(tlist(tint), tlist(tint)), _reverse),
+        Primitive("SORT", arrow(tlist(tint), tlist(tint)), _sort),
+        Primitive("SUM", arrow(tlist(tint), tint), _sum)
+        ] + [
+        Primitive("MAP", arrow(int_to_int, tlist(tint), tlist(tint)), _map), #is this okay???
+        Primitive("FILTER", arrow(int_to_bool, tlist(tint), tlist(tint)), _filter), #is this okay???
+        Primitive("COUNT", arrow(int_to_bool, tlist(tint), tint), _count), #is this okay???
+        Primitive("ZIPWITH", arrow(int_to_int_to_int, tlist(tint), tlist(tint), tlist(tint)), _zipwith), #is this okay???
+        Primitive("SCANL1", arrow(int_to_int_to_int, tlist(tint), tlist(tint)), _scanl1), #is this okay???
+        ] + [
+        Primitive("INC", int_to_int, _succ),
+        Primitive("DEC", int_to_int, _pred),
+        Primitive("SHL", int_to_int, _double),
+        Primitive("SHR", int_to_int, _half),
+        Primitive("doNEG", int_to_int, _negate),
+        Primitive("SQR", int_to_int, _square),
+        Primitive("MUL3", int_to_int, _triple),
+        Primitive("DIV3", int_to_int, _third),
+        Primitive("MUL4", int_to_int, _quad),
+        Primitive("DIV4", int_to_int, _quarter),
+        ] + [
+        Primitive("isPOS", int_to_bool, _pos),
+        Primitive("isNEG", int_to_bool, _neg),
+        Primitive("isEVEN", int_to_bool, _even),
+        Primitive("isODD", int_to_bool, _odd),
+        ] + [
+        Primitive("+", int_to_int_to_int, _add),
+        Primitive("-", int_to_int_to_int, _sub),
+        Primitive("*", int_to_int_to_int, _mult),
+        Primitive("MIN", int_to_int_to_int, _min),
+        Primitive("MAX", int_to_int_to_int, _max)       
+    ]
+
+def OldDeepcoderPrimitives():
+    return [
+        Primitive("head", arrow(tlist(tint), tint), _head), 
+        Primitive("tail", arrow(tlist(tint), tint), _tail),
+        Primitive("take", arrow(tint, tlist(tint), tlist(tint)), _take),
+        Primitive("drop", arrow(tint, tlist(tint), tlist(tint)), _drop),
+        Primitive("access", arrow(tint, tlist(tint), tint), _access),
+        Primitive("minimum", arrow(tlist(tint), tint), _minimum),
+        Primitive("maximum", arrow(tlist(tint), tint), _maximum),
+        Primitive("reverse", arrow(tlist(tint), tlist(tint)), _reverse),
+        Primitive("sort", arrow(tlist(tint), tlist(tint)), _sort),
+        Primitive("sum", arrow(tlist(tint), tint), _sum)
+        ] + [
+        Primitive("map", arrow(int_to_int, tlist(tint), tlist(tint)), _map), #is this okay???
+        Primitive("filter_int", arrow(int_to_bool, tlist(tint), tlist(tint)), _filter), #is this okay???
+        Primitive("count", arrow(int_to_bool, tlist(tint), tint), _count), #is this okay???
+        Primitive("zipwith", arrow(int_to_int_to_int, tlist(tint), tlist(tint), tlist(tint)), _zipwith), #is this okay???
+        Primitive("scanl1", arrow(int_to_int_to_int, tlist(tint), tlist(tint)), _scanl1), #is this okay???
+        # ] + [
+        # Primitive("succ", arrow(tint, tint), _succ),
+        # Primitive("pred", arrow(tint, tint), _pred),
+        # Primitive("double", arrow(tint, tint), _double),
+        # Primitive("half", arrow(tint, tint), _half),
+        # Primitive("neg", arrow(tint, tint), _neg),
+        # Primitive("square", arrow(tint, tint), _square),
+        # Primitive("triple", arrow(tint, tint), _triple),
+        # Primitive("third", arrow(tint, tint), _third),
+        # Primitive("quad", arrow(tint, tint), _quad),
+        # Primitive("quarter", arrow(tint, tint), _quarter),
+        # ] + [
+        # Primitive("pos", arrow(tint, tbool), _pos),
+        # Primitive("neg", arrow(tint, tbool), _neg),
+        # Primitive("even", arrow(tint, tbool), _even),
+        # Primitive("odd", arrow(tint, tbool), _odd),
+        # ] + [
+        # Primitive("add", arrow(tint, tint, tint), _add),
+        # Primitive("sub", arrow(tint, tint, tint), _sub),
+        # Primitive("mult", arrow(tint, tint, tint), _mult),
+        # Primitive("min", arrow(tint, tint, tint), _min),
+        # Primitive("max", arrow(tint, tint, tint), _max)
+        ] + [
+        Primitive("succ_fn", int_to_int, _succ),
+        Primitive("pred_fn", int_to_int, _pred),
+        Primitive("double_fn", int_to_int, _double),
+        Primitive("half_fn", int_to_int, _half),
+        Primitive("negate_fn", int_to_int, _negate),
+        Primitive("square_fn", int_to_int, _square),
+        Primitive("triple_fn", int_to_int, _triple),
+        Primitive("third_fn", int_to_int, _third),
+        Primitive("quad_fn", int_to_int, _quad),
+        Primitive("quarter_fn", int_to_int, _quarter),
+        ] + [
+        Primitive("pos_fn", int_to_bool, _pos),
+        Primitive("neg_fn", int_to_bool, _neg),
+        Primitive("even_fn", int_to_bool, _even),
+        Primitive("odd_fn", int_to_bool, _odd),
+        ] + [
+        Primitive("add_fn", int_to_int_to_int, _add),
+        Primitive("sub_fn", int_to_int_to_int, _sub),
+        Primitive("mult_fn", int_to_int_to_int, _mult),
+        Primitive("min_fn", int_to_int_to_int, _min),
+        Primitive("max_fn", int_to_int_to_int, _max)       
+    ]
+
+def deepcoderProductions():
+    return [(0.0, prim) for prim in deepcoderPrimitives()]
+
+def flatten_program(p):
+    string = p.show(False)
+    num_inputs = string.count('lambda')
+    string = string.replace('lambda', '')
+    string = string.replace('(', '')
+    string = string.replace(')', '')
+    #remove '_fn' (optional)
+    for i in range(num_inputs):
+        string = string.replace('$' + str(num_inputs-i-1),'input_' + str(i))
+    string = string.split(' ')
+    string = list(filter(lambda x: x is not '', string))
+    return string
+
+if __name__ == "__main__":
+    #g = Grammar.uniform(deepcoderPrimitives())
+    g = Grammar.fromProductions(deepcoderProductions(), logVariable=.9)
+    request = arrow(tlist(tint), tint, tint)
+    p = g.sample(request)
+    print("request:", request)
+    print("program:")
+    print(prettyProgram(p))
+    print("flattened_program:")
+    flat = flatten_program(p)
+    print(flat)
+
+    #robustfill output = names from productions + input_0-2 or 3
+
+
+
+
+
+    # # with open("/home/ellisk/om/ec/experimentOutputs/list_aic=1.0_arity=3_ET=1800_expandFrontier=2.0_it=4_likelihoodModel=all-or-nothing_MF=5_baseline=False_pc=10.0_L=1.0_K=5_rec=False.pickle", "rb") as handle:
+    # #     b = pickle.load(handle).grammars[-1]
+    # # print b
+
+    # p = Program.parse(
+    #     "(lambda (lambda (lambda (if (empty? $0) empty (cons (+ (car $1) (car $0)) ($2 (cdr $1) (cdr $0)))))))")
+    # t = arrow(tlist(tint), tlist(tint), tlist(tint))  # ,tlist(tbool))
+    # print(g.logLikelihood(arrow(t, t), p))
+    # assert False
+    # print(b.logLikelihood(arrow(t, t), p))
+
+    # # p = Program.parse("""(lambda (lambda
+    # # (unfold 0
+    # # (lambda (+ (index $0 $2) (index $0 $1)))
+    # # (lambda (1+ $0))
+    # # (lambda (eq? $0 (length $1))))))
+    # # """)
+    # p = Program.parse("""(lambda (lambda
+    # (map (lambda (+ (index $0 $2) (index $0 $1))) (range (length $0))  )))""")
+    # # .replace("unfold", "#(lambda (lambda (lambda (lambda (fix1 $0 (lambda (lambda (#(lambda (lambda (lambda (if $0 empty (cons $1 $2))))) ($1 ($3 $0)) ($4 $0) ($5 $0)))))))))").\
+    # # replace("length", "#(lambda (fix1 $0 (lambda (lambda (if (empty? $0) 0 (+ ($1 (cdr $0)) 1))))))").\
+    # # replace("forloop", "(#(lambda (lambda (lambda (lambda (fix1 $0 (lambda (lambda (#(lambda (lambda (lambda (if $0 empty (cons $1 $2))))) ($1 ($3 $0)) ($4 $0) ($5 $0))))))))) (lambda (#(eq? 0) $0)) $0 (lambda (#(lambda (- $0 1)) $0)))").\
+    # # replace("inc","#(lambda (+ $0 1))").\
+    # # replace("drop","#(lambda (lambda (fix2 $0 $1 (lambda (lambda (lambda (if
+    # # (#(eq? 0) $1) $0 (cdr ($2 (- $1 1) $0)))))))))"))
+    # print(p)
+    # print(g.logLikelihood(t, p))
+    # assert False
+
+    # print("??")
+    # p = Program.parse(
+    #     "#(lambda (#(lambda (lambda (lambda (fix1 $0 (lambda (lambda (if (empty? $0) $3 ($4 (car $0) ($1 (cdr $0)))))))))) (lambda $1) 1))")
+    # for j in range(10):
+    #     l = list(range(j))
+    #     print(l, p.evaluate([])(lambda x: x * 2)(l))
+    #     print()
+    # print()
+
+    # print("multiply")
+    # p = Program.parse(
+    #     "(lambda (lambda (lambda (if (eq? $0 0) 0 (+ $1 ($2 $1 (- $0 1)))))))")
+    # print(g.logLikelihood(arrow(arrow(tint, tint, tint), tint, tint, tint), p))
+    # print()
+
+    # print("take until 0")
+    # p = Program.parse("(lambda (lambda (if (eq? $1 0) empty (cons $1 $0))))")
+    # print(g.logLikelihood(arrow(tint, tlist(tint), tlist(tint)), p))
+    # print()
+
+    # print("countdown primitive")
+    # p = Program.parse(
+    #     "(lambda (lambda (if (eq? $0 0) empty (cons (+ $0 1) ($1 (- $0 1))))))")
+    # print(
+    #     g.logLikelihood(
+    #         arrow(
+    #             arrow(
+    #                 tint, tlist(tint)), arrow(
+    #                 tint, tlist(tint))), p))
+    # print(_fix(9)(p.evaluate([])))
+    # print("countdown w/ better primitives")
+    # p = Program.parse(
+    #     "(lambda (lambda (if (eq0 $0) empty (cons (+1 $0) ($1 (-1 $0))))))")
+    # print(
+    #     g.logLikelihood(
+    #         arrow(
+    #             arrow(
+    #                 tint, tlist(tint)), arrow(
+    #                 tint, tlist(tint))), p))
+
+    # print()
+
+    # print("prepend zeros")
+    # p = Program.parse(
+    #     "(lambda (lambda (lambda (if (eq? $1 0) $0 (cons 0 ($2 (- $1 1) $0))))))")
+    # print(
+    #     g.logLikelihood(
+    #         arrow(
+    #             arrow(
+    #                 tint,
+    #                 tlist(tint),
+    #                 tlist(tint)),
+    #             tint,
+    #             tlist(tint),
+    #             tlist(tint)),
+    #         p))
+    # print()
+    # assert False
+
+    # p = Program.parse(
+    #     "(lambda (fix1 $0 (lambda (lambda (if (empty? $0) 0 (+ 1 ($1 (cdr $0))))))))")
+    # print(p.evaluate([])(list(range(17))))
+    # print(g.logLikelihood(arrow(tlist(tbool), tint), p))
+
+    # p = Program.parse(
+    #     "(lambda (lambda (if (empty? $0) 0 (+ 1 ($1 (cdr $0))))))")
+    # print(
+    #     g.logLikelihood(
+    #         arrow(
+    #             arrow(
+    #                 tlist(tbool), tint), arrow(
+    #                 tlist(tbool), tint)), p))
+
+    # p = Program.parse(
+    #     "(lambda (fix1 $0 (lambda (lambda (if (empty? $0) 0 (+ (car $0) ($1 (cdr $0))))))))")
+
+    # print(p.evaluate([])(list(range(4))))
+    # print(g.logLikelihood(arrow(tlist(tint), tint), p))
+
+    # p = Program.parse(
+    #     "(lambda (lambda (if (empty? $0) 0 (+ (car $0) ($1 (cdr $0))))))")
+    # print(p)
+    # print(
+    #     g.logLikelihood(
+    #         arrow(
+    #             arrow(
+    #                 tlist(tint),
+    #                 tint),
+    #             tlist(tint),
+    #             tint),
+    #         p))
+
+    # print("take")
+    # p = Program.parse(
+    #     "(lambda (lambda (lambda (if (eq? $1 0) empty (cons (car $0) ($2 (- $1 1) (cdr $0)))))))")
+    # print(p)
+    # print(
+    #     g.logLikelihood(
+    #         arrow(
+    #             arrow(
+    #                 tint,
+    #                 tlist(tint),
+    #                 tlist(tint)),
+    #             tint,
+    #             tlist(tint),
+    #             tlist(tint)),
+    #         p))
+    # assert False
+
+    # print(p.evaluate([])(list(range(4))))
+    # print(g.logLikelihood(arrow(tlist(tint), tlist(tint)), p))
+
+    # p = Program.parse(
+    #     """(lambda (fix (lambda (lambda (match $0 0 (lambda (lambda (+ $1 ($3 $0))))))) $0))""")
+    # print(p.evaluate([])(list(range(4))))
+    # print(g.logLikelihood(arrow(tlist(tint), tint), p))

dreamcoder/domains/misc/napsPrimitives.py

@@ -0,0 +1,198 @@
+#napsPrimitives.py
+from dreamcoder.program import Primitive, prettyProgram
+from dreamcoder.grammar import Grammar
+from dreamcoder.type import tlist, tint, arrow, baseType #, t0, t1, t2
+
+#from functools import reduce
+
+
+#types
+PROGRAM = baseType("PROGRAM")
+
+RECORD = baseType("RECORD")
+FUNC = baseType("FUNC")
+
+VAR = baseType("VAR")
+STMT = baseType("STMT")
+EXPR = baseType("EXPR")
+ASSIGN = baseType("ASSIGN")
+LHS = baseType("LHS")
+IF = baseType("IF")
+FOREACH = baseType("FOREACH")
+WHILE = baseType("WHILE")
+BREAK = baseType("BREAK")
+CONTINUE = baseType("CONTINUE")
+RETURN = baseType("RETURN")
+NOOP = baseType("NOOP")
+FIELD = baseType("FIELD")
+CONSTANT = baseType("CONSTANT")
+INVOKE = baseType("INVOKE")
+TERNARY = baseType("TERNARY")
+CAST = baseType("CAST")
+TYPE = baseType("TYPE")
+
+#other types
+function_name = baseType("function_name")
+field_name = baseType("field_name")
+name = baseType("name")  # for records and functions
+value = baseType("value")
+
+# definitions:
+
+def _program(records): return lambda funcs: {'types': records, 'funcs': funcs}
+# record
+def _func(string): return lambda tp: lambda name: lambda vars1: lambda vars2: lambda stmts: [string, tp, name, vars1, vars2, stmts]
+def _var(tp): return lambda name: ['var', tp, name]
+# stmt
+# expr
+def _assign(tp): return lambda lhs: lambda expr: ['assign', tp, lhs, expr]
+# lhs 
+def _if(tp): return lambda expr: lambda stmts1: lambda stmts2: ['if', tp, expr, stmts1, stmts2]  # TODO
+def _foreach(tp): return lambda var: lambda expr: lambda stmts: ['foreach', tp, expr, stmts]  # TODO
+def _while(tp): return lambda expr: lambda stmts1: lambda stmts1: ['while', tp, expr, stmts1, stmts2]  # or: ['while', tp, expr, [stmts1], [stmts2]] #TODO
+# break
+# continue
+def _return(tp): return lambda expr: ['return', tp, expr]
+# noop
+def _field(tp): return lambda expr: lambda field_name: ['field', tp, expr, field_name]
+def _constant(tp): return lambda value: ['val', tp, value]  #TODO deal with value
+def _invoke(tp): return lambda function_name: lambda exprs: ['invoke', tp, function_name, exprs]  # TODO, deal with fn_name and lists
+def _ternary(tp): return lambda expr1: lambda expr2: lambda expr3: ['?:', tp, expr1, expr2, expr3]
+def _cast(tp): return lambda expr: ['cast', tp, expr]
+
+# types:
+
+# TODO: deal with lists - x 
+# TODO: deal with names
+# TODO: deal with values - x 
+
+# TODO: deal with the program/record __main__ and __globals__ stuff
+
+
+
+def napsPrimitives():
+    return [
+        Primitive("program", arrow(tlist(RECORD), tlist(FUNC), PROGRAM), _program), # TODO
+        # RECORD
+        Primitive("func", arrow(TYPE, name, tlist(VAR), tlist(VAR), tlist(VAR), tlist(STMT)), _func('func')), # TODO
+        Primitive("ctor", arrow(TYPE, name, tlist(VAR), tlist(VAR), tlist(VAR), tlist(STMT)), _func('ctor')),
+        Primitive("var", arrow(TYPE, name, VAR), _var)
+        ] + [ 
+        # STMT ::= EXPR | IF | FOREACH | WHILE | BREAK | CONTINUE | RETURN | NOOP
+        Primitive("stmt_expr", arrow(EXPR, STMT), lambda x: x),
+        Primitive("stmt_if", arrow(IF, STMT), lambda x: x),
+        Primitive("stmt_foreach", arrow(FOREACH, STMT), lambda x: x),
+        Primitive("stmt_while", arrow(WHILE, STMT), lambda x: x),
+        Primitive("stmt_break", arrow(BREAK, STMT), lambda x: x),
+        Primitive("stmt_continue", arrow(CONTINUE, STMT), lambda x: x),
+        Primitive("stmt_return", arrow(RETURN, STMT), lambda x: x),
+        Primitive("stmt_noop", arrow(NOOP, STMT), lambda x: x)
+        ] + [
+        # EXPR ::= ASSIGN | VAR | FIELD | CONSTANT | INVOKE | TERNARY | CAST
+        Primitive("expr_assign", arrow(ASSIGN, EXPR), lambda x: x),
+        Primitive("expr_var", arrow(VAR, EXPR), lambda x: x),
+        Primitive("expr_field", arrow(FIELD, EXPR), lambda x: x),
+        Primitive("expr_constant", arrow(CONSTANT, EXPR), lambda x: x),
+        Primitive("expr_invoke", arrow(INVOKE, EXPR), lambda x: x),
+        Primitive("expr_ternary", arrow(TERNARY, EXPR), lambda x: x),
+        Primitive("expr_cast", arrow(CAST, EXPR), lambda x: x)
+        ] + [
+        Primitive("assign", arrow(TYPE, LHS, EXPR, ASSIGN), _assign)
+        ] + [
+        # LHS ::= VAR | FIELD | INVOKE
+        Primitive("lhs_var", arrow(VAR, LHS), lambda x: x),
+        Primitive("lhs_field", arrow(FIELD, LHS), lambda x: x),
+        Primitive("lhs_invoke", arrow(INVOKE, LHS), lambda x: x)
+        ] + [
+        Primitive("if", arrow(TYPE, EXPR, tlist(STMT), tlist(STMT), IF), _if),
+        Primitive("foreach", arrow(TYPE, VAR, EXPR, tlist(STMT), FOREACH), _foreach),
+        Primitive("while", arrow(TYPE, EXPR, tlist(STMT), tlist(STMT), WHILE), _while),
+        Primitive("break", arrow(TYPE, BREAK), lambda tp: ['break', tp]),
+        Primitive("continue", arrow(TYPE, CONTINUE), lambda tp: ['continue', tp]),
+        Primitive("return", arrow(TYPE, EXPR, RETURN), _return),
+        Primitive("noop", NOOP, ['noop']),
+        Primitive("field", arrow(TYPE, EXPR, field_name, FIELD), _field),  # TODO
+        Primitive("constant", arrow(TYPE, value, CONSTANT), _constant),
+        Primitive("invoke", arrow(TYPE, function_name, tlist(EXPR), INVOKE), _invoke),  # TODO
+        Primitive("ternary", arrow(TYPE, EXPR, EXPR, EXPR, TERNARY), _ternary),
+        Primitive("cast", arrow(TYPE, EXPR, CAST), _cast)
+        ] + [
+        # below are TYPE:
+        Primitive("bool", TYPE, 'bool'),
+        Primitive("char", TYPE, 'char'),
+        Primitive("char*", TYPE, 'char*'),
+        Primitive("int", TYPE, 'int'),
+        Primitive("real", TYPE, 'real'),
+        Primitive("array", arrow(TYPE, TYPE), lambda tp: tp + '*'),
+        Primitive("set", arrow(TYPE, TYPE), lambda tp: tp + '%'),
+        Primitive("map", arrow(TYPE, TYPE, TYPE), lambda tp1: lambda tp2: '<'+tp1+'|'+tp2+'>'),
+        Primitive("record_name", TYPE, 'record_name#')  # TODO
+        ] + [
+        #stuff about lists:
+        # STMTs, EXPRs, VARs, maybe Funcs and records
+        Primitive('list_init_stmt', arrow(STMT, tlist(STMT)), lambda stmt: [stmt]),
+        Primitive('list_add_stmt', arrow(STMT, tlist(STMT), tlist(STMT)), lambda stmt: lambda stmts: stmts + [stmt]),
+        Primitive('list_init_expr', arrow(EXPR, tlist(EXPR)), lambda expr: [expr]),
+        Primitive('list_add_expr', arrow(EXPR, tlist(EXPR), tlist(EXPR)), lambda expr: lambda exprs: exprs + [expr]),
+        Primitive('list_init_var', arrow(VAR, tlist(VAR)), lambda var: [var]),
+        Primitive('list_add_var', arrow(VAR, tlist(VAR), tlist(VAR)), lambda var: lambda _vars: _vars + [var])
+        ] + [
+        # value
+        Primitive('0', value, 0),
+        Primitive("1", value, "1"),
+        Primitive("-1", value, "-1")
+        # ...  
+        ] + [
+        # function_name:
+        Primitive('+', function_name, '+'),
+        Primitive('&&', function_name, "&&"),
+        Primitive("!", function_name, "!"),
+        Primitive("!=", function_name, "!="),
+        Primitive("string_find", function_name,"string_find")
+        # ... 
+        ] + [
+        # field_name:
+        Primitive('', field_name, '')
+        # ...
+        ] + [
+        # 
+        Primitive(f'var{str(i)}', name, f'var{str(i)}') for i in range(12)
+    ]
+
+
+#for first pass, can just hard code vars and maps n stuff
+
+def ec_prog_to_uast(prog):  # TODO
+    # ideally, just evaluate and then parse
+    uast = prog.evaluate([])
+    return uast
+
+def deepcoderProductions():
+    return [(0.0, prim) for prim in deepcoderPrimitives()]
+
+# def flatten_program(p):
+#     string = p.show(False)
+#     num_inputs = string.count('lambda')
+#     string = string.replace('lambda', '')
+#     string = string.replace('(', '')
+#     string = string.replace(')', '')
+#     #remove '_fn' (optional)
+#     for i in range(num_inputs):
+#         string = string.replace('$' + str(num_inputs-i-1),'input_' + str(i))
+#     string = string.split(' ')
+#     string = list(filter(lambda x: x is not '', string))
+#     return string
+
+if __name__ == "__main__":
+    #g = Grammar.uniform(deepcoderPrimitives())
+    g = Grammar.fromProductions(deepcoderProductions(), logVariable=.9)
+    request = arrow(tlist(tint), tint, tint)
+    p = g.sample(request)
+    print("request:", request)
+    print("program:")
+    print(prettyProgram(p))
+    print("flattened_program:")
+    flat = flatten_program(p)
+    print(flat)
+
+  

dreamcoder/domains/regex/groundtruthRegexes.py

@@ -0,0 +1,172 @@
+
+#dict of gt regexes
+
+"""
+ 		pre.create(".+"),
+ 		pre.create("\d+"),
+        pre.create("\w+"),
+        pre.create("\s+"),
+        pre.create("\\u+"),
+        pre.create("\l+")
+"""
+gt_dict = {
+    776: "JPC\\u\\u\\d+\\.png",
+    922: "WHS\\d_\\d+",
+    354: "\\u+",
+    523: "(\\u)+|\\.",
+    184: "\\.\\d+",
+    501: "u\\d\\d",
+    760: "\\u\\u",
+    49: "(\\u)+\\u\\d?",
+    732: "\\uR5\\d\\d",
+    450: "-\\d(\\.(\\d)+)?",
+    350: "\\u\\u",
+    467: "hu\\d(\\d|\\u)+",
+    622: "A(\\d|\\u)**",
+    476: "\\u+",
+    554: "\\u\\u",
+    940: "\\u\\u?",
+    496: "\\u\\u",
+    369: "\\u\\u\\u",
+    596: "\\u+",
+    720: "\\(\\d\\d\\d\\) \\d\\d\\d-\\d\\d\\d\\d",
+    53: "rec-\\d\\d\\d?-(org)|(dup-0)",
+    150: "N\\d\\d",
+    741: "#\\d\\d\\d",
+    18: "A|C-\\d+-\\d+",
+    589: "A(\\u|\\d)++",
+    666: "\\(\\d\\d\\d\\) \\d\\d\\d-\\d\\d\\d\\d",
+    581: "us13\\u\\d\\d",
+    299: "E07000\\d\\d\\d",
+    638: "\\l+\\d+\\l+\\d+",
+    364: "\\u\\u",
+    334: "-00:\\d\\d:\\d\\d.\\d",
+    38: "SRX89\\d+",
+    247: "'\\d\\d:\\d\\d:00'",
+    506: "(S|H)\\d+",
+    891: "(r|v)\\d?",
+    911: "KW-\\d+",
+    792: "\\d*\\u*",
+    508: "N000\\d+",
+    842: "-?\\d?\\d\\.\\d\\d%",
+    200: "\\u\\u",
+    694: "\\(\\d+\\)",
+    210: "(\\d(\\.\\d)?)|(--)",
+    298: "DS_25(\\u|\\d)+",
+    668: "\\u+",
+    939: "ms0\\d+",
+    944: "\\u+\\d?",
+    731: "ManH.0\\d\\d",
+    229: "\\u+(-\\u+)?",
+    28: "Y201\\d/\\d\\d\\d\\d",
+    374: "q000\\d(_000\\d)?",
+    819: "\\d*\\l*\\d*",
+    516: "-122.3\\d+",
+    417: "\\u\\uT\\uB",
+    660: "ENGL?\\d\\d\\d",
+    585: "M?\\u+",
+    325: "BUS M \\d\\d\\d.*",
+    823: "\\u\\u\\u",
+    515: "L|\\u - (\\?\\?)|(\\d?\\d\\.\\d lbs\\.)",
+    864: "\\u+",
+    359: "MAM\\.OSBS\\.201\\d\\.\\d\\d",
+    594: "(\\u|\\d)+( (\\u|\\d)+)*",
+    788: "-\\d(,\\d+)?",
+    188: "cat\\. \\d\\d",
+    355: ".+",
+    799: "\\u\\d\\d",
+    902: "\\u\\d\\d",
+    920: "A\\.\\d\\d",
+    330: "Resp\\d\\d",
+    396: "\\u+(( |/)\\u+)?",
+    393: "US $ \\d\\.\\d\\d",
+    680: "Z:-?0\\.\\d\\d",
+    744: "t1_cv(\\l|\\d)+",
+    461: "(\\u|\\l)+\\d+",
+    631: "$\\d+\\.\\d+",
+    195: "(OLE)?\\d+",
+    693: "\\u",
+    577: "EFO_000\\d+",
+    392: "$\\d+(,\\d\\d\\d)*\\.00",
+    688: "\\u+( \\u+)*",
+    816: "\\u\\u\\u",
+    489: "UK\\u\\d",
+    251: "\\l\\l\\l",
+    653: "C\\d+",
+    769: "(\\u|\\l|\\d|-)+\\d+",
+    991: "Q\\d-201\\d",
+    342: "\\u\\u\\d\\d\\d\\d",
+    308: "\\u\\u\\u\\u",
+    136: "IMPC_\\u\\u\\u_\\d\\d\\d_\\d\\d\\d",
+    327: "#\\d+((/|-)\\d+)*",
+    981: "\\u\\u\\u",
+    892: "(.|\\l)*",
+    375: "P\\u\\.\\d\\d\\d\\d\\.\\d\\d\\d",
+    499: "A000\\d+",
+    474: "\\u+",
+    50: "V06\\d+",
+    381: "F?\\d+",
+    883: "-79.\\d+",
+    173: "(\\u|\\l)+\\d+",
+    147: "\\u\\u\\u-\\u\\u\\u",
+    419: "\\u\\u",
+    961: "-?\\d\\.\\d*",
+    148: "Q\\d\\d",
+    975: "(\\d|\\u)+",
+    79: "\\d+(,\\d\\d\\d)+",
+    775: "\\u\\l\\l \\d+ \\d\\d\\d\\d",
+    774: "FOS\\d\\d+",
+    561: ".+",
+    509: "S000\\d+",
+    494: "S1900\\d+",
+    119: "$\\d\\d(,\\d\\d\\d)+",
+    29: "(\\u|\\l|\\d)+",
+    121: "(\\d|\\u|\\.|/|\\(|\\))+",
+    61: "R \\d\\d\\d.\\d\\d",
+    871: "-0.7\\d+",
+    639: "\\u+?\\d+",
+    729: "COMISARIA \\d\\d",
+    193: "\\u\\d\\d",
+    752: "(.*|\\u\\.?)+",
+    17: "$\\d.\\d\\d",
+    914: "R\\d\\d\\d\\d",
+    510: "P\\d000\\d\\d\\d\\d",
+    443: "(W|L) \\d-\\d+",
+    20: "MDEL\\d\\d?\\.\\d\\l",
+    64: "c04p0100(\\l|\\d)",
+    301: "(\\u|\\d)+(-(\\u|\\d)+)*",
+    664: "N\\d",
+    493: "[0\\.0\\d+]",
+    765: "-?\\d\\.\\d+( \\(0\\.\\d+\\))?"
+
+
+}
+badRegexTasks = {
+    "Data column no. 922",
+    "Data column no. 184",
+    "Data column no. 467",
+    "Data column no. 476",
+    "Data column no. 150",
+    "Data column no. 299",
+    "Data column no. 334",
+    "Data column no. 493",
+    "Data column no. 891",
+    "Data column no. 792",
+    "Data column no. 765",
+    "Data column no. 944",
+    "Data column no. 374",
+    "Data column no. 660",
+    "Data column no. 188",
+    "Data column no. 920",
+    "Data column no. 330",
+    "Data column no. 396",
+    "Data column no. 680",
+    "Data column no. 769",
+    "Data column no. 308",
+    "Data column no. 375",
+    "Data column no. 474",
+    "Data column no. 79",
+    "Data column no. 871",
+    "Data column no. 729",
+    "Data column no. 664",
+}

dreamcoder/domains/regex/main.py

@@ -0,0 +1,384 @@
+# analog of list.py for regex tasks. Responsible for actually running the task.
+
+from dreamcoder.domains.regex.makeRegexTasks import makeOldTasks, makeLongTasks, makeShortTasks, makeWordTasks, makeNumberTasks, makeHandPickedTasks, makeNewTasks, makeNewNumberTasks
+from dreamcoder.domains.regex.regexPrimitives import basePrimitives, altPrimitives, easyWordsPrimitives, alt2Primitives, concatPrimitives, reducedConcatPrimitives, strConstConcatPrimitives, PRC
+from dreamcoder.dreamcoder import explorationCompression, Task
+from dreamcoder.grammar import Grammar
+from dreamcoder.likelihoodModel import add_cutoff_values, add_string_constants
+from dreamcoder.program import Abstraction, Application
+from dreamcoder.type import tpregex
+from dreamcoder.utilities import eprint, flatten, testTrainSplit, POSITIVEINFINITY
+
+import random
+import math
+import pregex as pre
+import os
+
+try:
+    from dreamcoder.recognition import RecurrentFeatureExtractor, JSONFeatureExtractor
+    class LearnedFeatureExtractor(RecurrentFeatureExtractor):
+        H = 64
+        special = 'regex'
+
+        def tokenize(self, examples):
+            def sanitize(l): return [z if z in self.lexicon else "?"
+                                     for z_ in l
+                                     for z in (z_ if isinstance(z_, list) else [z_])]
+
+            tokenized = []
+            for xs, y in examples:
+                if isinstance(y, list):
+                    y = ["LIST_START"] + y + ["LIST_END"]
+                else:
+                    y = [y]
+                y = sanitize(y)
+                if len(y) > self.maximumLength:
+                    return None
+
+                serializedInputs = []
+                for xi, x in enumerate(xs):
+                    if isinstance(x, list):
+                        x = ["LIST_START"] + x + ["LIST_END"]
+                    else:
+                        x = [x]
+                    x = sanitize(x)
+                    if len(x) > self.maximumLength:
+                        return None
+                    serializedInputs.append(x)
+
+                tokenized.append((tuple(serializedInputs), y))
+
+            return tokenized
+
+        def __init__(self, tasks, testingTasks=[], cuda=False):
+            self.lexicon = set(flatten((t.examples for t in tasks + testingTasks), abort=lambda x: isinstance(
+                x, str))).union({"LIST_START", "LIST_END", "?"})
+
+            self.num_examples_list = [len(t.examples) for t in tasks]
+
+            # Calculate the maximum length
+            self.maximumLength = POSITIVEINFINITY
+            self.maximumLength = max(len(l)
+                                     for t in tasks + testingTasks
+                                     for xs, y in self.tokenize(t.examples)
+                                     for l in [y] + [x for x in xs])
+
+            super(
+                LearnedFeatureExtractor,
+                self).__init__(
+                lexicon=list(
+                    self.lexicon),
+                tasks=tasks,
+                cuda=cuda,
+                H=self.H,
+                bidirectional=True)
+            self.parallelTaskOfProgram = False
+
+
+        def taskOfProgram(self, p, t):
+            #raise NotImplementedError
+            num_examples = random.choice(self.num_examples_list)
+
+            p = p.visit(ConstantInstantiateVisitor.SINGLE)
+
+            preg = p.evaluate([])(pre.String(""))
+            t = Task("Helm", t, [((), list(preg.sample())) for _ in range(num_examples) ])
+            return t
+except: pass
+        #in init: loop over tasks, save lengths, 
+
+
+class ConstantInstantiateVisitor(object):
+    def __init__(self):
+        self.regexes = [
+        pre.create(".+"),
+        pre.create("\d+"),
+        pre.create("\w+"),
+        pre.create("\s+"),
+        pre.create("\\u+"),
+        pre.create("\l+")]
+
+    def primitive(self, e):
+        if e.name == "r_const":
+            #return Primitive("STRING", e.tp, random.choice(self.words))
+            s = random.choice(self.regexes).sample() #random string const
+            s = pre.String(s)
+            e.value = PRC(s,arity=0)
+        return e
+
+    def invented(self, e): return e.body.visit(self)
+
+    def index(self, e): return e
+
+    def application(self, e):
+        return Application(e.f.visit(self), e.x.visit(self))
+
+    def abstraction(self, e):
+        return Abstraction(e.body.visit(self))
+#TODO fix
+
+
+
+
+class MyJSONFeatureExtractor(JSONFeatureExtractor):
+    N_EXAMPLES = 5
+
+    def _featuresOfProgram(self, program, tp):
+        try:
+            preg = program.evaluate([])
+            # if 'left_paren' in program.show(False):
+            #eprint("string_pregex:", string_pregex)
+            #eprint("string_pregex:", string_pregex)
+
+        except IndexError:
+            # free variable
+            return None
+        except Exception as e:
+            eprint("Exception during evaluation:", e)
+            if "Attempt to evaluate fragment variable" in e:
+                eprint("program (bc fragment error)", program)
+            return None
+
+        examples = []
+
+        for _ in range(self.N_EXAMPLES * 5):  # oh this is arbitrary ig
+
+            try:
+                y = preg.sample()  # TODO
+
+                #this line should keep inputs short, so that helmholtzbatch can be large
+                #allows it to try other samples
+                #(Could also return None off the bat... idk which is better)
+                #if len(y) > 20:
+                #    continue
+                #eprint(tp, program, x, y)
+                examples.append(y)
+            except BaseException:
+                continues
+            if len(examples) >= self.N_EXAMPLES:
+                break
+        else:
+            return None
+        return examples  # changed to list_features(examples) from examples
+
+
+def regex_options(parser):
+    parser.add_argument("--maxTasks", type=int,
+                        default=500,
+                        help="truncate tasks to fit within this boundary")
+    parser.add_argument(
+        "--maxExamples",
+        type=int,
+        default=10,
+        help="truncate number of examples per task to fit within this boundary")
+    parser.add_argument("--tasks",
+                        default="long",
+                        help="which tasks to use",
+                        choices=["old", "short", "long", "words", "number", "handpicked", "new", "newNumber"])
+    parser.add_argument("--primitives",
+                        default="concat",
+                        help="Which primitive set to use",
+                        choices=["base", "alt1", "easyWords", "alt2", "concat", "reduced", "strConst"])
+    parser.add_argument("--extractor", type=str,
+                        choices=["hand", "deep", "learned", "json"],
+                        default="learned")  # if i switch to json it breaks
+    parser.add_argument("--split", metavar="TRAIN_RATIO",
+                        type=float,
+                        default=0.8,
+                        help="split test/train")
+    parser.add_argument("-H", "--hidden", type=int,
+                        default=256,
+                        help="number of hidden units")
+    parser.add_argument("--likelihoodModel",
+                        default="probabilistic",
+                        help="likelihood Model",
+                        choices=["probabilistic", "all-or-nothing"])
+    parser.add_argument("--topk_use_map",
+                        dest="topk_use_only_likelihood",
+                        action="store_false")
+    parser.add_argument("--debug",
+                        dest="debug",
+                        action="store_true")
+    parser.add_argument("--ll_cutoff",
+                        dest="use_ll_cutoff",
+                        nargs='*',
+                        default=False,
+                        help="use ll cutoff for training tasks (for probabilistic likelihood model only). default is False,")
+    parser.add_argument("--use_str_const",
+                        action="store_true",
+                        help="use string constants")
+
+    """parser.add_argument("--stardecay",
+                        type=float,
+                        dest="stardecay",
+                        default=0.5,
+                        help="p value for kleenestar and plus")"""
+
+# Lucas recommends putting a struct with the definitions of the primitives here.
+# TODO:
+# Build likelihood funciton
+# modify NN
+# make primitives
+# make tasks
+
+
+def main(args):
+    """
+    Takes the return value of the `commandlineArguments()` function as input and
+    trains/tests the model on regular expressions.
+    """
+    #for dreaming
+
+    #parse use_ll_cutoff
+    use_ll_cutoff = args.pop('use_ll_cutoff')
+    if not use_ll_cutoff is False:
+
+        #if use_ll_cutoff is a list of strings, then train_ll_cutoff and train_ll_cutoff 
+        #will be tuples of that string followed by the actual model
+
+        if len(use_ll_cutoff) == 1:
+            train_ll_cutoff = use_ll_cutoff[0] # make_cutoff_model(use_ll_cutoff[0], tasks))
+            test_ll_cutoff = use_ll_cutoff[0] # make_cutoff_model(use_ll_cutoff[0], tasks))
+        else:
+            assert len(use_ll_cutoff) == 2
+            train_ll_cutoff = use_ll_cutoff[0] #make_cutoff_model(use_ll_cutoff[0], tasks))
+            test_ll_cutoff = use_ll_cutoff[1] #make_cutoff_model(use_ll_cutoff[1], tasks))
+    else:
+        train_ll_cutoff = None
+        test_ll_cutoff = None
+
+
+    regexTasks = {"old": makeOldTasks,
+                "short": makeShortTasks,
+                "long": makeLongTasks,
+                "words": makeWordTasks,
+                "number": makeNumberTasks,
+                "handpicked": makeHandPickedTasks,
+                "new": makeNewTasks,
+                "newNumber": makeNewNumberTasks
+                }[args.pop("tasks")]
+
+    tasks = regexTasks()  # TODO
+    eprint("Generated", len(tasks), "tasks")
+
+    maxTasks = args.pop("maxTasks")
+    if len(tasks) > maxTasks:
+        eprint("Unwilling to handle {} tasks, truncating..".format(len(tasks)))
+        seed = 42 # previously this was hardcoded and never changed
+        random.seed(seed)
+        random.shuffle(tasks)
+        del tasks[maxTasks:]
+
+    maxExamples = args.pop("maxExamples")
+   
+
+    split = args.pop("split")
+    test, train = testTrainSplit(tasks, split)
+    eprint("Split tasks into %d/%d test/train" % (len(test), len(train)))
+
+
+    test = add_cutoff_values(test, test_ll_cutoff)
+    train = add_cutoff_values(train, train_ll_cutoff)
+    eprint("added cutoff values to tasks, train: ", train_ll_cutoff, ", test:", test_ll_cutoff )
+
+
+    if args.pop("use_str_const"):
+        assert args["primitives"] == "strConst" or args["primitives"] == "reduced"
+        ConstantInstantiateVisitor.SINGLE = \
+            ConstantInstantiateVisitor()
+        test = add_string_constants(test)
+        train = add_string_constants(train)
+        eprint("added string constants to test and train")
+    
+    for task in test + train:
+        if len(task.examples) > maxExamples:
+            task.examples = task.examples[:maxExamples]
+
+        task.specialTask = ("regex", {"cutoff": task.ll_cutoff, "str_const": task.str_const})
+        task.examples = [(xs, [y for y in ys ])
+                         for xs,ys in task.examples ]
+        task.maxParameters = 1
+
+    # from list stuff
+    primtype = args.pop("primitives")
+    prims = {"base": basePrimitives,
+             "alt1": altPrimitives,
+             "alt2": alt2Primitives,
+             "easyWords": easyWordsPrimitives,
+             "concat": concatPrimitives,
+             "reduced": reducedConcatPrimitives,
+             "strConst": strConstConcatPrimitives
+             }[primtype]
+
+    extractor = {
+        "learned": LearnedFeatureExtractor,
+        "json": MyJSONFeatureExtractor
+    }[args.pop("extractor")]
+
+    extractor.H = args.pop("hidden")
+
+    #stardecay = args.stardecay
+    #stardecay = args.pop('stardecay')
+    #decaystr = 'd' + str(stardecay)
+    import datetime
+
+    timestamp = datetime.datetime.now().isoformat()
+    outputDirectory = "experimentOutputs/regex/%s"%timestamp
+    os.system("mkdir -p %s"%outputDirectory)
+
+    args.update({
+        "featureExtractor": extractor,
+        "outputPrefix": "%s/regex"%(outputDirectory),
+        "evaluationTimeout": 0.005,
+        "topk_use_only_likelihood": True,
+        "maximumFrontier": 10,
+        "compressor": args.get("compressor","ocaml")
+    })
+    ####
+
+
+        # use the
+    #prim_list = prims(stardecay)
+    prim_list = prims()
+    specials = ["r_kleene", "r_plus", "r_maybe", "r_alt", "r_concat"]
+    n_base_prim = len(prim_list) - len(specials)
+
+    productions = [
+        (math.log(0.5 / float(n_base_prim)),
+         prim) if prim.name not in specials else (
+            math.log(0.10),
+            prim) for prim in prim_list]
+
+
+    baseGrammar = Grammar.fromProductions(productions, continuationType=tpregex)
+    #baseGrammar = Grammar.uniform(prims())
+
+    #for i in range(100):
+    #    eprint(baseGrammar.sample(tpregex))
+
+    #eprint(baseGrammar)
+    #explore
+    test_stuff = args.pop("debug")
+    if test_stuff:
+        eprint(baseGrammar)
+        eprint("sampled programs from prior:")
+        for i in range(100): #100
+            eprint(baseGrammar.sample(test[0].request,maximumDepth=1000))
+        eprint("""half the probability mass is on higher-order primitives.
+Therefore half of enumerated programs should have more than one node.
+However, we do not observe this.
+Instead we see a very small fraction of programs have more than one node. 
+So something seems to be wrong with grammar.sample.
+
+Furthermore: observe the large print statement above. 
+This prints the candidates for sampleDistribution in grammar.sample.
+the first element of each tuple is the probability passed into sampleDistribution.
+Half of the probability mass should be on the functions, but instead they are equally 
+weighted with the constants. If you look at the grammar above, this is an error!!!!
+""")
+        assert False
+
+    del args["likelihoodModel"]
+    explorationCompression(baseGrammar, train,
+                           testingTasks = test,
+                           **args)

dreamcoder/domains/regex/makeRegexTasks.py

@@ -0,0 +1,347 @@
+import dill
+import os
+import json
+from string import printable
+
+import sys
+try:
+    from pregex import pregex
+except:
+    print("Failure to load pregex. This is only acceptable if using pypy",file=sys.stderr)
+
+from dreamcoder.task import Task
+from dreamcoder.type import tpregex, arrow
+from dreamcoder.utilities import get_data_dir
+
+
+def makeOldTasks():
+    # a series of tasks
+
+    taskfile = os.path.join(get_data_dir(), 'data_filtered.json')
+    #task_list = pickle.load(open(taskfile, 'rb'))
+
+    with open(taskfile) as f:
+        file_contents = f.read()
+    task_list = json.loads(file_contents)
+
+
+    # if I were to just dump all of them:
+    regextasks = [
+        Task("Luke data column no." + str(i),
+             arrow(tpregex, tpregex),
+                 [((), example) for example in task_list[i]]
+             ) for i in range(len(task_list))]
+
+
+    """ regextasks = [
+        Task("length bool", arrow(none,tstr),
+             [((l,), len(l))
+              for _ in range(10)
+              for l in [[flip() for _ in range(randint(0,10)) ]] ]),
+        Task("length int", arrow(none,tstr),
+             [((l,), len(l))
+              for _ in range(10)
+              for l in [randomList()] ]),
+    ]
+  """
+    return regextasks  # some list of tasks
+
+
+
+
+def makeShortTasks():
+
+    #load new data:
+
+    taskfile = os.path.join(get_data_dir(), "regex_data_csv_900.p")
+
+    with open(taskfile, 'rb') as handle:
+        data = dill.load(handle)
+
+    tasklist = data[0][:100] #a list of indices
+
+    regextasks = [
+        Task("Data column no. " + str(i),
+            arrow(tpregex, tpregex),
+            [((), example) for example in task] 
+        ) for i, task in enumerate(tasklist)]
+
+
+
+    return regextasks
+
+def makeLongTasks():
+
+    #load new data:
+
+    taskfile = os.path.join(get_data_dir(), "regex_data_csv_900.p")
+
+    with open(taskfile, 'rb') as handle:
+        data = dill.load(handle)
+
+    tasklist = data[0] #a list of indices
+
+    regextasks = [
+        Task("Data column no. " + str(i),
+            arrow(tpregex, tpregex),
+            [((), example) for example in task] 
+        ) for i, task in enumerate(tasklist)]
+
+
+
+    return regextasks
+
+def makeWordTasks():
+
+    #load new data:
+
+    taskfile = os.path.join(get_data_dir(), "regex_data_csv_900.p")
+
+    with open(taskfile, 'rb') as handle:
+        data = dill.load(handle)
+
+    tasklist = data[0] #a list of indices
+
+
+
+
+    all_upper = [0, 2, 8, 9, 10, 11, 12, 17, 18, 19, 20, 22]
+    all_lower = [1]
+
+    # match_col(data[0],'\\u(\l+)')
+    one_capital_lower_plus = [144, 200, 241, 242, 247, 296, 390, 392, 444, 445, 481, 483, 485, 489, 493, 542, 549, 550, 581]
+
+    #match_col(data[0],'(\l ?)+')
+    lower_with_maybe_spaces = [1, 42, 47, 99, 100, 102, 201, 246, 248, 293, 294, 345, 437, 545, 590]
+
+    #match_col(data[0],'(\\u\l+ ?)+')
+    capital_then_lower_maybe_spaces = [144, 200, 241, 242, 247, 296, 390, 392, 395, 438, 444, 445, 481, 483, 484, 485, 487, 489, 493, 494, 542, 546, 549, 550, 578, 581, 582, 588, 591, 624, 629]
+
+    #match_col(data[0],'(\\u+ ?)+')
+    all_caps_spaces = [0, 2, 8, 9, 10, 11, 12, 17, 18, 19, 20, 22, 25, 26, 35, 36, 43, 45, 46, 49, 50, 52, 56, 59, 87, 89, 95, 101, 140, 147, 148, 149, 199, 332, 336, 397, 491, 492, 495, 580, 610]
+
+    #one_capital_and_lower = [566, 550, 549, 542, 505, 493, 494, 489, 488, 485, 483, 481, 445, 444, 438, 296, 241, 242, 200, ]
+    #all_lower_with_a_space = [545]
+    #all_lower_maybe_space = [534]
+    #one_capital_lower_maybe_spaces = [259, 262, 263, 264]
+
+
+    #full_list = test_list + train_list
+    train_list = []
+    full_list = all_upper + all_lower + one_capital_lower_plus + lower_with_maybe_spaces + capital_then_lower_maybe_spaces + all_caps_spaces
+
+    regextasks = [
+        Task("Data column no. " + str(i),
+            arrow(tpregex, tpregex),
+            [((), example) for example in task] 
+        ) for i, task in enumerate(tasklist) if i in full_list ]
+
+    for i in train_list:
+        regextasks[i].mustTrain = True
+
+
+    return regextasks
+
+def makeNumberTasks():
+
+    #load new data:
+
+    taskfile = os.path.join(get_data_dir(), "regex_data_csv_900.p")
+
+    with open(taskfile, 'rb') as handle:
+        data = dill.load(handle)
+
+    tasklist = data[0] #a list of indices
+
+
+
+    #match_col(data[0],'\d*\.\d*')
+    raw_decimals = [121, 122, 163, 164, 165, 170, 172, 173, 175, 178, 218, 228, 230, 231, 252, 253,
+    254, 258, 259, 305, 320, 330, 334, 340, 348, 350, 351, 352, 353, 355, 357, 358, 361, 363, 364, 
+    371, 380, 382, 409, 410, 411, 447, 448, 449, 450, 458, 469, 471, 533, 562, 564]
+
+
+    decimals_pos_neg_dollar = [3, 4, 5, 6, 7, 13, 16, 24, 27, 28, 29, 30, 31, 32, 33, 34, 37, 38, 39, 40, 
+    53, 54, 55, 57, 58, 60, 61, 63, 64, 65, 66, 68, 69, 70, 71, 73, 74, 77, 78, 80, 81, 103, 104, 105, 
+    106, 107, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 121, 122, 123, 124, 125, 126, 128,
+     129, 131, 132, 134, 135, 139, 146, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165,
+      166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 180, 181, 182, 183, 184, 185, 186,
+       193, 194, 195, 204, 205, 207, 209, 210, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 223, 224,
+        225, 226, 227, 228, 229, 230, 231, 232, 249, 250, 251, 252, 253, 254, 255, 256, 258, 259, 260, 261,
+         263, 266, 267, 270, 271, 272, 277, 299, 301, 302, 305, 306, 307, 309, 312, 313, 315, 319, 320, 324,
+          326, 327, 330, 334, 340, 348, 350, 351, 352, 353, 354, 355, 356, 357, 358, 361, 362, 363, 364, 368,
+           371, 373, 377, 380, 382, 400, 401, 402, 403, 405, 406, 409, 410, 411, 413, 435, 439, 446, 447, 448,
+            449, 450, 451, 452, 453, 456, 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 469, 470, 471, 477,
+             498, 500, 502, 503, 507, 512, 518, 519, 520, 532, 533, 553, 554, 555, 556, 557, 558, 559, 560, 561,
+              562, 564, 565, 572, 577]
+
+    #match_col(data[0],'(\d*,?\d*)+')
+    commas = []
+    #match_col(data[0],'(\d*,?\d*)+')
+    commas_and_all = []
+
+    #full_list = test_list + train_list
+    train_list = []
+    full_list = decimals_pos_neg_dollar
+
+    regextasks = [
+        Task("Data column no. " + str(i),
+            arrow(tpregex, tpregex),
+            [((), example) for example in task] 
+        ) for i, task in enumerate(tasklist) if i in full_list ]
+
+    for i in train_list:
+        regextasks[i].mustTrain = True
+
+
+    return regextasks
+
+
+def makeHandPickedTasks():
+
+    #load new data:
+
+    taskfile = os.path.join(get_data_dir(), "regex_data_csv_900.p")
+
+    with open(taskfile, 'rb') as handle:
+        data = dill.load(handle)
+
+    tasklist = data[0] #a list of indices
+
+
+    full_list = list(range(199)) + \
+    [209,218,222,223,224,225,226] + \
+    list(range(222,233)) + \
+    [235,237,238,239,243,244,245,252,253,254,255,257,258,259,260,261,264,265,269,272,274] + \
+    list(range(275,291)) + \
+    [295,297,300,303,304,305,306,310,311,312,314,315,316,320,321,323,327,329,330,333,334,335,337,338,339,340,341,342,343,344] + \
+    list(range(348,359)) + \
+    [361,369,373,379,380,382,387,403,405,407,408] + \
+    list(range(409,417)) + \
+    list(range(418,437)) + \
+    list(range(440,444)) + \
+    list(range(446,452)) + \
+    list(range(456,460)) + \
+    list(range(466,472)) + \
+    [503,504]
+
+
+    regextasks = [
+        Task("Data column no. " + str(i),
+            arrow(tpregex, tpregex),
+            [((), example) for example in task] 
+        ) for i, task in enumerate(tasklist) if i in full_list ]
+
+    #for i in train_list:
+    #    regextasks[i].mustTrain = True
+
+
+    return regextasks
+
+def makeNewTasks(include_only=None):
+
+    #load new data:
+
+    taskfile = os.path.join(get_data_dir(), "csv_filtered_all_background_novel.p")
+
+    with open(taskfile, 'rb') as handle:
+        data = dill.load(handle)
+
+    tasklist = data['background'] #a list of indices
+
+    if include_only:
+        regextasks = [
+            Task("Data column no. " + str(i),
+                arrow(tpregex, tpregex),
+                [((), example) for example in task['train']] 
+            ) for i, task in enumerate(tasklist) if i in include_only]
+    else:
+        regextasks = [
+            Task("Data column no. " + str(i),
+                arrow(tpregex, tpregex),
+                [((), example) for example in task['train']] 
+            ) for i, task in enumerate(tasklist)]
+
+    #for i in train_list:
+    #    regextasks[i].mustTrain = True
+
+    return regextasks
+REGEXTASKS = None
+def regexHeldOutExamples(task, include_only=None):
+
+    #load new data:
+    global REGEXTASKS
+    if REGEXTASKS is None:
+        taskfile = os.path.join(get_data_dir(), "csv_filtered_all_background_novel.p")
+
+        with open(taskfile, 'rb') as handle:
+            data = dill.load(handle)
+
+        tasklist = data['background'] #a list of indices
+
+        if include_only:
+            regextasks = [
+                Task("Data column no. " + str(i),
+                    arrow(tpregex, tpregex),
+                    [((), example) for example in _task['test']] 
+                ) for i, _task in enumerate(tasklist) if i in include_only]
+        else:
+            regextasks = [
+                Task("Data column no. " + str(i),
+                    arrow(tpregex, tpregex),
+                    [((), example) for example in _task['test']] 
+                ) for i, _task in enumerate(tasklist)]
+
+        #for i in train_list:
+        #    regextasks[i].mustTrain = True
+
+        REGEXTASKS = {t.name: t.examples for t in regextasks}
+    fullTask = REGEXTASKS[task.name]
+    return fullTask
+        
+
+
+def makeNewNumberTasks():
+
+    tasks = makeNewTasks()
+    numberTasks = [t for t in tasks if not any(p in ex for p in printable[10:62] for _, ex in t.examples)]
+    return numberTasks 
+
+
+# a helper function which takes a list of lists and sees which match a specific regex.
+def match_col(dataset, rstring):
+    r = pregex.create(rstring)
+    matches = []
+    for i, col in enumerate(dataset):
+        score = sum([r.match(example) for example in col])
+        if score != float('-inf'):
+            matches.append(i)
+    return matches
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--include_only",
+                        default=None,
+                        nargs="+",
+                        type=int)
+    args = parser.parse_args()
+
+
+    def show_tasks(dataset):
+        task_list = []
+        for task in dataset:
+            print(task.name)
+            print([example[1] for example in task.examples[:20]])
+            task_list.append([example[1] for example in task.examples])
+        return task_list
+
+    task = {"number": makeNumberTasks,
+    "words": makeWordTasks,
+    "all": makeLongTasks,
+    "new": makeNewTasks}['new']
+
+
+    x = show_tasks(task(args.include_only))
+
+    

dreamcoder/domains/regex/regexPrimitives.py

@@ -0,0 +1,367 @@
+import sys
+from dreamcoder.program import Primitive
+from dreamcoder.grammar import Grammar
+from dreamcoder.type import arrow, tpregex
+from string import printable
+
+try:
+    from pregex import pregex
+except:
+    print("Failure to load pregex. This is only acceptable if using pypy",file=sys.stderr)
+
+
+# evaluation to regular regex form. then I can unflatten using Luke's stuff.
+
+
+def _kleene(x): return pregex.KleeneStar(x, p=0.25)
+
+
+def _plus(x): return pregex.Plus(x, p=0.25)
+
+
+def _maybe(x): return pregex.Maybe(x)
+
+
+# maybe should be reversed#"(" + x + "|" + y + ")"
+def _alt(x): return lambda y: pregex.Alt([x, y])
+
+
+def _concat(x): return lambda y: pregex.Concat([x, y])  # "(" + x + y + ")"
+
+
+#For sketch:
+def _kleene_5(x): return pregex.KleeneStar(x)
+
+def _plus_5(x): return pregex.Plus(x)
+
+
+disallowed = [
+    ("#", "hash"),
+    ("!", "bang"),
+    ("\"", "double_quote"),
+    ("$", "dollar"),
+    ("%", "percent"),
+    ("&", "ampersand"),
+    ("'", "single_quote"),
+    (")", "left_paren"),
+    ("(", "right_paren"),
+    ("*", "astrisk"),
+    ("+", "plus"),
+    (",", "comma"),
+    ("-", "dash"),
+    (".", "period"),
+    ("/", "slash"),
+    (":", "colon"),
+    (";", "semicolon"),
+    ("<", "less_than"),
+    ("=", "equal"),
+    (">", "greater_than"),
+    ("?", "question_mark"),
+    ("@", "at"),
+    ("[", "left_bracket"),
+    ("\\", "backslash"),
+    ("]", "right_bracket"),
+    ("^", "carrot"),
+    ("_", "underscore"),
+    ("`", "backtick"),
+    ("|", "bar"),
+    ("}", "right_brace"),
+    ("{", "left_brace"),
+    ("~", "tilde"),
+    (" ", "space"),
+    ("\t", "tab")
+]
+
+disallowed_list = [char for char, _ in disallowed]
+
+class PRC(): #PregexContinuation
+    def __init__(self, f, arity=0, args=[]):
+        self.f = f  
+        self.arity = arity
+        self.args = args
+
+    def __call__(self, pre):
+
+        if self.arity == len(self.args):
+            if self.arity == 0: return pregex.Concat([self.f, pre]) 
+            elif self.arity == 1: return pregex.Concat([self.f(*self.args), pre])
+            else: return pregex.Concat([self.f(self.args), pre]) #this line is bad, need brackets around input to f if f is Alt
+        else: return PRC(self.f, self.arity, args=self.args+[pre(pregex.String(""))])
+
+
+def concatPrimitives():
+    return [Primitive("string_" + i, arrow(tpregex, tpregex), PRC(pregex.String(i))) for i in printable[:-4] if i not in disallowed_list
+            ] + [
+        Primitive("string_" + name, arrow(tpregex, tpregex), PRC(pregex.String(char))) for char, name in disallowed
+    ] + [
+        Primitive("r_dot", arrow(tpregex, tpregex), PRC(pregex.dot)),
+        Primitive("r_d", arrow(tpregex, tpregex), PRC(pregex.d)),
+        Primitive("r_s", arrow(tpregex, tpregex), PRC(pregex.s)),
+        Primitive("r_w", arrow(tpregex, tpregex), PRC(pregex.w)),
+        Primitive("r_l", arrow(tpregex, tpregex), PRC(pregex.l)),
+        Primitive("r_u", arrow(tpregex, tpregex), PRC(pregex.u)),
+        #todo
+        Primitive("r_kleene", arrow(arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.KleeneStar,1)),
+        Primitive("r_plus", arrow(arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.Plus,1)),
+        Primitive("r_maybe", arrow(arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.Maybe,1)),
+        Primitive("r_alt", arrow(arrow(tpregex, tpregex) , arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.Alt, 2)),
+    ]
+
+def strConstConcatPrimitives():
+    return [Primitive("string_" + i, arrow(tpregex, tpregex), PRC(pregex.String(i))) for i in printable[:-4] if i not in disallowed_list
+            ] + [
+        Primitive("string_" + name, arrow(tpregex, tpregex), PRC(pregex.String(char))) for char, name in disallowed
+    ] + [
+        Primitive("r_dot", arrow(tpregex, tpregex), PRC(pregex.dot)),
+        Primitive("r_d", arrow(tpregex, tpregex), PRC(pregex.d)),
+        Primitive("r_s", arrow(tpregex, tpregex), PRC(pregex.s)),
+        Primitive("r_w", arrow(tpregex, tpregex), PRC(pregex.w)),
+        Primitive("r_l", arrow(tpregex, tpregex), PRC(pregex.l)),
+        Primitive("r_u", arrow(tpregex, tpregex), PRC(pregex.u)),
+        #todo
+        Primitive("r_kleene", arrow(arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.KleeneStar,1)),
+        Primitive("r_plus", arrow(arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.Plus,1)),
+        Primitive("r_maybe", arrow(arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.Maybe,1)),
+        Primitive("r_alt", arrow(arrow(tpregex, tpregex) , arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.Alt, 2)),
+    ] + [
+    Primitive("r_const", arrow(tpregex, tpregex), None)
+    ]
+
+
+def reducedConcatPrimitives():
+    #uses strConcat!!
+    #[Primitive("empty_string", arrow(tpregex, tpregex), PRC(pregex.String("")))
+            #] + [
+    return [Primitive("string_" + i, arrow(tpregex, tpregex), PRC(pregex.String(i))) for i in printable[:-4] if i not in disallowed_list
+            ] + [
+        Primitive("string_" + name, arrow(tpregex, tpregex), PRC(pregex.String(char))) for char, name in disallowed
+        ] + [
+        Primitive("r_dot", arrow(tpregex, tpregex), PRC(pregex.dot)),
+        Primitive("r_d", arrow(tpregex, tpregex), PRC(pregex.d)),
+        Primitive("r_s", arrow(tpregex, tpregex), PRC(pregex.s)),
+        #Primitive("r_w", arrow(tpregex, tpregex), PRC(pregex.w)),
+        Primitive("r_l", arrow(tpregex, tpregex), PRC(pregex.l)),
+        Primitive("r_u", arrow(tpregex, tpregex), PRC(pregex.u)),
+        #todo
+        Primitive("r_kleene", arrow(arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.KleeneStar,1)),
+        #Primitive("r_plus", arrow(arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.Plus,1)),
+        Primitive("r_maybe", arrow(arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.Maybe,1)),
+        Primitive("r_alt", arrow(arrow(tpregex, tpregex) , arrow(tpregex, tpregex), arrow(tpregex,tpregex)), PRC(pregex.Alt, 2)),
+    ] + [
+    Primitive("r_const", arrow(tpregex, tpregex), None)
+    ]
+
+
+def sketchPrimitives():
+    return [Primitive("string_" + i, tpregex, pregex.String(i)) for i in printable[:-4] if i not in disallowed_list
+            ] + [
+        Primitive("string_" + name, tpregex, pregex.String(char)) for char, name in disallowed
+    ] + [
+        Primitive("r_dot", tpregex, pregex.dot),
+        Primitive("r_d", tpregex, pregex.d),
+        Primitive("r_s", tpregex, pregex.s),
+        Primitive("r_w", tpregex, pregex.w),
+        Primitive("r_l", tpregex, pregex.l),
+        Primitive("r_u", tpregex, pregex.u),
+        Primitive("r_kleene", arrow(tpregex, tpregex), _kleene_5),
+        Primitive("r_plus", arrow(tpregex, tpregex), _plus_5),
+        Primitive("r_maybe", arrow(tpregex, tpregex), _maybe),
+        Primitive("r_alt", arrow(tpregex, tpregex, tpregex), _alt),
+        Primitive("r_concat", arrow(tpregex, tpregex, tpregex), _concat),
+    ]
+
+def basePrimitives():
+    return [Primitive("string_" + i, tpregex, pregex.String(i)) for i in printable[:-4] if i not in disallowed_list
+            ] + [
+        Primitive("string_" + name, tpregex, pregex.String(char)) for char, name in disallowed
+    ] + [
+        Primitive("r_dot", tpregex, pregex.dot),
+        Primitive("r_d", tpregex, pregex.d),
+        Primitive("r_s", tpregex, pregex.s),
+        Primitive("r_w", tpregex, pregex.w),
+        Primitive("r_l", tpregex, pregex.l),
+        Primitive("r_u", tpregex, pregex.u),
+        Primitive("r_kleene", arrow(tpregex, tpregex), _kleene),
+        Primitive("r_plus", arrow(tpregex, tpregex), _plus),
+        Primitive("r_maybe", arrow(tpregex, tpregex), _maybe),
+        Primitive("r_alt", arrow(tpregex, tpregex, tpregex), _alt),
+        Primitive("r_concat", arrow(tpregex, tpregex, tpregex), _concat),
+    ]
+
+
+
+def altPrimitives():
+    return [
+        Primitive("empty_string", tpregex, pregex.String(""))
+    ] + [
+        Primitive("string_" + i, tpregex, pregex.String(i)) for i in printable[:-4] if i not in disallowed_list
+    ] + [
+        Primitive("string_" + name, tpregex, pregex.String(char)) for char, name in disallowed
+    ] + [
+        Primitive("r_dot", tpregex, pregex.dot),
+        Primitive("r_d", tpregex, pregex.d),
+        Primitive("r_s", tpregex, pregex.s),
+        Primitive("r_w", tpregex, pregex.w),
+        Primitive("r_l", tpregex, pregex.l),
+        Primitive("r_u", tpregex, pregex.u),
+        Primitive("r_kleene", arrow(tpregex, tpregex), _kleene),
+        #Primitive("r_plus", arrow(tpregex, tpregex), _plus),
+        Primitive("r_maybe", arrow(tpregex, tpregex), _maybe),
+        Primitive("r_alt", arrow(tpregex, tpregex, tpregex), _alt),
+        Primitive("r_concat", arrow(tpregex, tpregex, tpregex), _concat),
+    ]
+
+def alt2Primitives():
+    return [
+        Primitive("empty_string", tpregex, pregex.String(""))
+    ] + [
+        Primitive("string_" + i, tpregex, pregex.String(i)) for i in printable[:-4] if i not in disallowed_list
+    ] + [
+        Primitive("string_" + name, tpregex, pregex.String(char)) for char, name in disallowed
+    ] + [
+        Primitive("r_dot", tpregex, pregex.dot),
+        Primitive("r_d", tpregex, pregex.d),
+        Primitive("r_s", tpregex, pregex.s),
+        Primitive("r_w", tpregex, pregex.w),
+        Primitive("r_l", tpregex, pregex.l),
+        Primitive("r_u", tpregex, pregex.u),
+        Primitive("r_kleene", arrow(tpregex, tpregex), _kleene),
+        #Primitive("r_plus", arrow(tpregex, tpregex), _plus),
+        #Primitive("r_maybe", arrow(tpregex, tpregex), _maybe),
+        Primitive("r_alt", arrow(tpregex, tpregex, tpregex), _alt),
+        Primitive("r_concat", arrow(tpregex, tpregex, tpregex), _concat),
+    ]
+
+def easyWordsPrimitives():
+    return [
+        Primitive("string_" + i, tpregex, pregex.String(i)) for i in printable[10:62] if i not in disallowed_list
+    ] + [
+        Primitive("r_d", tpregex, pregex.d),
+        Primitive("r_s", tpregex, pregex.s),
+        #Primitive("r_w", tpregex, pregex.w),
+        Primitive("r_l", tpregex, pregex.l),
+        Primitive("r_u", tpregex, pregex.u),
+        Primitive("r_kleene", arrow(tpregex, tpregex), _kleene),
+        Primitive("r_plus", arrow(tpregex, tpregex), _plus),
+        Primitive("r_maybe", arrow(tpregex, tpregex), _maybe),
+        Primitive("r_alt", arrow(tpregex, tpregex, tpregex), _alt),
+        Primitive("r_concat", arrow(tpregex, tpregex, tpregex), _concat),
+    ]
+
+
+#def _wrapper(x): return lambda y: y
+
+#specials = [".","*","+","?","|"]
+"""
+>>> import pregex as pre
+>>> abc = pre.CharacterClass("abc", [0.1, 0.1, 0.8], name="MyConcept")
+>>> abc.sample()
+'b'
+>>> abc.sample()
+'c'
+>>> abc.sample()
+'c'
+>>> abc.match("c")
+-0.2231435513142097
+>>> abc.match("a")
+-2.3025850929940455
+>>> abc
+MyConcept
+>>> x = pre.KleeneStar(abc)
+>>> x.match("aabbac")
+-16.58809928020405
+>>> x.sample()
+''
+>>> x.sample()
+''
+>>> x.sample()
+'cbcacc'
+>>> x
+(KleeneStar 0.5 MyConcept)
+>>> str(x)
+'MyConcept*'
+"""
+
+
+def emp_dot(corpus): return pregex.CharacterClass(printable[:-4], emp_distro_from_corpus(corpus, printable[:-4]), name=".")
+
+def emp_d(corpus): return pregex.CharacterClass(printable[:10], emp_distro_from_corpus(corpus, printable[:10]), name="\\d")
+
+#emp_s = pre.CharacterClass(slist, [], name="emp\\s") #may want to forgo this one. 
+
+def emp_dot_no_letter(corpus): return pregex.CharacterClass(printable[:10]+printable[62:], emp_distro_from_corpus(corpus, printable[:10]+printable[62:]), name=".")
+
+def emp_w(corpus): return pregex.CharacterClass(printable[:62], emp_distro_from_corpus(corpus, printable[:62]), name="\\w")
+
+def emp_l(corpus): return pregex.CharacterClass(printable[10:36], emp_distro_from_corpus(corpus, printable[10:36]), name="\\l")
+
+def emp_u(corpus): return pregex.CharacterClass(printable[36:62], emp_distro_from_corpus(corpus, printable[36:62]), name="\\u")
+
+
+def emp_distro_from_corpus(corpus, char_list):
+    from collections import Counter
+    c = Counter(char for task in corpus for example in task.examples for string in example[1] for char in string)
+    n = sum(c[char] for char in char_list)
+    return [c[char]/n for char in char_list]
+
+
+
+def matchEmpericalPrimitives(corpus):
+    return lambda: [
+        Primitive("empty_string", tpregex, pregex.String(""))
+    ] + [
+        Primitive("string_" + i, tpregex, pregex.String(i)) for i in printable[:-4] if i not in disallowed_list
+    ] + [
+        Primitive("string_" + name, tpregex, pregex.String(char)) for char, name in disallowed
+    ] + [
+        Primitive("r_dot", tpregex, emp_dot(corpus) ),
+        Primitive("r_d", tpregex, emp_d(corpus) ),
+        Primitive("r_s", tpregex, pregex.s),
+        Primitive("r_w", tpregex, emp_w(corpus) ),
+        Primitive("r_l", tpregex, emp_l(corpus) ),
+        Primitive("r_u", tpregex, emp_u(corpus) ),
+        Primitive("r_kleene", arrow(tpregex, tpregex), _kleene),
+        #Primitive("r_plus", arrow(tpregex, tpregex), _plus),
+        Primitive("r_maybe", arrow(tpregex, tpregex), _maybe),
+        Primitive("r_alt", arrow(tpregex, tpregex, tpregex), _alt),
+        Primitive("r_concat", arrow(tpregex, tpregex, tpregex), _concat),
+    ]
+
+def matchEmpericalNoLetterPrimitives(corpus):
+    return lambda: [
+        Primitive("empty_string", tpregex, pregex.String(""))
+    ] + [
+        Primitive("string_" + i, tpregex, pregex.String(i)) for i in printable[:-4] if i not in disallowed_list + list(printable[10:62])
+    ] + [
+        Primitive("string_" + name, tpregex, pregex.String(char)) for char, name in disallowed
+    ] + [
+        Primitive("r_dot", tpregex, emp_dot_no_letter(corpus) ),
+        Primitive("r_d", tpregex, emp_d(corpus) ),
+        Primitive("r_s", tpregex, pregex.s),
+        Primitive("r_kleene", arrow(tpregex, tpregex), _kleene),
+        #Primitive("r_plus", arrow(tpregex, tpregex), _plus),
+        #Primitive("r_maybe", arrow(tpregex, tpregex), _maybe),
+        Primitive("r_alt", arrow(tpregex, tpregex, tpregex), _alt),
+        Primitive("r_concat", arrow(tpregex, tpregex, tpregex), _concat),
+    ]
+
+
+if __name__=='__main__':
+    concatPrimitives()
+    from dreamcoder.program import Program
+
+    p=Program.parse("(lambda (r_kleene (lambda (r_maybe (lambda (string_x $0)) $0)) $0))")
+    print(p)
+    print(p.runWithArguments([pregex.String("")]))
+
+    prims = concatPrimitives()
+    g = Grammar.uniform(prims)
+
+    for i in range(100):
+        prog = g.sample(arrow(tpregex,tpregex))
+        preg = prog.runWithArguments([pregex.String("")])
+        print("preg:", preg.__repr__())
+        print("sample:", preg.sample())
+
+
+

dreamcoder/domains/text/main.py

@@ -0,0 +1,270 @@
+from dreamcoder.dreamcoder import ecIterator
+from dreamcoder.domains.text.makeTextTasks import makeTasks, loadPBETasks
+from dreamcoder.domains.text.textPrimitives import primitives
+from dreamcoder.domains.list.listPrimitives import bootstrapTarget
+from dreamcoder.enumeration import *
+
+import os
+import datetime
+import random
+from functools import reduce
+import dill
+
+
+class ConstantInstantiateVisitor(object):
+    def __init__(self, words):
+        self.words = words
+
+    def primitive(self, e):
+        if e.name == "STRING":
+            return Primitive("STRING", e.tp, random.choice(self.words))
+        return e
+
+    def invented(self, e): return e.body.visit(self)
+
+    def index(self, e): return e
+
+    def application(self, e):
+        return Application(e.f.visit(self), e.x.visit(self))
+
+    def abstraction(self, e):
+        return Abstraction(e.body.visit(self))
+
+
+try:
+    from dreamcoder.recognition import *
+    class LearnedFeatureExtractor(RecurrentFeatureExtractor):
+        special = 'string'
+
+        def tokenize(self, examples):
+            def tokenize_example(xs,y):
+                if not isinstance(y, list): y = [y]
+                return xs,y
+            return [tokenize_example(*e) for e in examples]
+
+        def __init__(self, tasks, testingTasks=[], cuda=False):
+            lexicon = {c
+                       for t in tasks + testingTasks
+                       for xs, y in self.tokenize(t.examples)
+                       for c in reduce(lambda u, v: u + v, list(xs) + [y])}
+            self.recomputeTasks = True
+
+            super(LearnedFeatureExtractor, self).__init__(lexicon=list(lexicon),
+                                                          H=64,
+                                                          tasks=tasks,
+                                                          bidirectional=True,
+                                                          cuda=cuda)
+            self.MAXINPUTS = 8
+
+        def taskOfProgram(self, p, tp):
+            # Instantiate STRING w/ random words
+            p = p.visit(ConstantInstantiateVisitor.SINGLE)
+            return super(LearnedFeatureExtractor, self).taskOfProgram(p, tp)
+except:
+    pass
+
+### COMPETITION CODE
+
+def competeOnOneTask(checkpoint, task,
+                     CPUs=8, timeout=3600, evaluationTimeout=0.0005):
+    if checkpoint.recognitionModel is not None:
+        recognizer = checkpoint.recognitionModel
+        challengeFrontiers, times, bestSearchTime = \
+                recognizer.enumerateFrontiers([task], 
+                                              CPUs=CPUs,
+                                              maximumFrontier=1,
+                                              enumerationTimeout=timeout,
+                                              evaluationTimeout=evaluationTimeout)
+    else:
+        challengeFrontiers, times, bestSearchTimes = \
+                multicoreEnumeration(checkpoint.grammars[-1], [task], 
+                                     CPUs=CPUs,
+                                     maximumFrontier=1,
+                                     enumerationTimeout=timeout,
+                                     evaluationTimeout=evaluationTimeout)
+    if len(times) == 0: return None, task
+    assert len(times) == 1
+    return times[0], task
+
+        
+
+def sygusCompetition(checkpoints, tasks):
+    from pathos.multiprocessing import Pool
+    import datetime
+    
+    # map from task to list of search times, one for each checkpoint.
+    # search time will be None if it is not solved
+    searchTimes = {t: [] for t in tasks}
+
+    CPUs = int(8/len(checkpoints))
+    maxWorkers = int(numberOfCPUs()/CPUs)
+    workers = Pool(maxWorkers)
+    eprint(f"You gave me {len(checkpoints)} checkpoints to ensemble. Each checkpoint will get {CPUs} CPUs. Creating a pool of {maxWorkers} worker processes.")
+    timeout = 3600
+
+
+    promises = []
+    for t in tasks:
+        for checkpoint in checkpoints:
+            promise = workers.apply_async(competeOnOneTask,
+                                          (checkpoint,t),
+                                          {"CPUs": CPUs,
+                                           "timeout": timeout})
+            promises.append(promise)
+    eprint(f"Queued {len(promises)} jobs.")
+    for promise in promises:
+        dt, task = promise.get()
+        if dt is not None:
+            searchTimes[task].append(dt)
+
+    searchTimes = {t: min(ts) if len(ts) > 0 else None
+                   for t,ts in searchTimes.items()}
+    
+    fn = "experimentOutputs/text_competition_%s.p"%(datetime.datetime.now().isoformat())
+    with open(fn,"wb") as handle:
+        pickle.dump(searchTimes, handle)
+    eprint()
+
+    hits = sum( t is not None for t in searchTimes.values() )
+    total = len(searchTimes)
+    percentage = 100*hits/total
+    eprint("Hits %d/%d = %f\n"%(hits, total, percentage))
+    eprint()
+    eprint("Exported competition results to",fn)
+    
+    
+
+def text_options(parser):
+    parser.add_argument(
+        "--showTasks",
+        action="store_true",
+        default=False,
+        help="show the training test and challenge tasks and then exit")
+    parser.add_argument(
+        "--trainChallenge",
+        action="store_true",
+        default=False,
+        help="Incorporate a random 50% of the challenge problems into the training set")
+    parser.add_argument(
+        "--onlyChallenge",
+        action="store_true",
+        default=False,
+        help="Only train on challenge problems and have testing problems.")
+    parser.add_argument(
+        "--latest",
+        action="store_true",
+        default=False,
+        help="evaluate on latest sygus problems rather than problems used in ec2 paper")
+    parser.add_argument(
+        "--noMap", action="store_true", default=False,
+        help="Disable built-in map primitive")
+    parser.add_argument(
+        "--noLength", action="store_true", default=False,
+        help="Disable built-in length primitive")
+    parser.add_argument(
+        "--noUnfold", action="store_true", default=False,
+        help="Disable built-in unfold primitive")
+    parser.add_argument(
+        "--compete",
+        nargs='+',
+        default=None,
+        type=str,
+        help="Do a simulated sygus competition (1hr+8cpus/problem) on the sygus tasks, restoring from provided checkpoint(s). If multiple checkpoints are provided, then we ensemble the models.")
+
+
+def main(arguments):
+    """
+    Takes the return value of the `commandlineArguments()` function as input and
+    trains/tests the model on manipulating sequences of text.
+    """
+
+    tasks = makeTasks()
+    eprint("Generated", len(tasks), "tasks")
+
+    for t in tasks:
+        t.mustTrain = False
+
+    test, train = testTrainSplit(tasks, 1.)
+    eprint("Split tasks into %d/%d test/train" % (len(test), len(train)))
+
+    latest = arguments.pop("latest")
+    challenge, challengeCheating = loadPBETasks("data/sygus" if latest else "PBE_Strings_Track")
+    eprint("Got %d challenge PBE tasks" % len(challenge))
+
+    if arguments.pop('trainChallenge'):
+        challengeTest, challengeTrain = testTrainSplit(challenge, 0.5)
+        challenge = challengeTest
+        train += challengeTrain
+        eprint(
+            "Incorporating %d (50%%) challenge problems into the training set." %
+            (len(challengeTrain)),
+            "We will evaluate on the held out challenge problems.",
+            "This makes a total of %d training problems." %
+            len(train))
+
+    if arguments.pop('onlyChallenge'):
+        train = challenge
+        test = []
+        challenge = []
+        eprint("Training only on sygus problems.")
+        
+
+    ConstantInstantiateVisitor.SINGLE = \
+        ConstantInstantiateVisitor(list(map(list, list({tuple([c for c in s])
+                                                        for t in test + train + challenge
+                                                        for s in t.stringConstants}))))
+
+    haveLength = not arguments.pop("noLength")
+    haveMap = not arguments.pop("noMap")
+    haveUnfold = not arguments.pop("noUnfold")
+    eprint(f"Including map as a primitive? {haveMap}")
+    eprint(f"Including length as a primitive? {haveLength}")
+    eprint(f"Including unfold as a primitive? {haveUnfold}")
+    baseGrammar = Grammar.uniform(primitives + [p
+                                                for p in bootstrapTarget()
+                                                if (p.name != "map" or haveMap) and \
+                                                (p.name != "unfold" or haveUnfold) and \
+                                                (p.name != "length" or haveLength)])
+    challengeGrammar = baseGrammar  # Grammar.uniform(targetTextPrimitives)
+
+    evaluationTimeout = 0.0005
+    # We will spend 10 minutes on each challenge problem
+    challengeTimeout = 10 * 60
+
+    for t in train + test + challenge:
+        t.maxParameters = 2
+
+    if arguments.pop("showTasks"):
+        for source, ts in [("train",tasks),("test",test),("challenge",challenge)]:
+            print(source,"tasks:")
+            for t in ts:
+                print(t.name)
+                for xs, y in t.examples:
+                    xs = ['"' + "".join(x) + '"' for x in xs]
+                    y = "".join(y) if isinstance(y,list) else y
+                    print('f(%s) = "%s"' % (", ".join(xs), y))
+                print("\t{%s}" % (t.stringConstants))
+            print()
+        sys.exit(0)
+
+
+    competitionCheckpoints = arguments.pop("compete")
+    if competitionCheckpoints:
+        checkpoints = []
+        for competitionCheckpoint in competitionCheckpoints:
+            with open(competitionCheckpoint, 'rb') as handle:
+                checkpoints.append(dill.load(handle))
+        sygusCompetition(checkpoints, challenge)
+        sys.exit(0)
+
+    timestamp = datetime.datetime.now().isoformat()
+    outputDirectory = "experimentOutputs/text/%s"%timestamp
+    os.system("mkdir -p %s"%outputDirectory)
+
+    generator = ecIterator(baseGrammar, train,
+                           testingTasks=test + challenge,
+                           outputPrefix="%s/text"%outputDirectory,
+                           evaluationTimeout=evaluationTimeout,
+                           **arguments)
+    for result in generator:
+        pass

dreamcoder/domains/text/makeTextTasks.py

@@ -0,0 +1,424 @@
+from dreamcoder.task import *
+from dreamcoder.type import *
+from dreamcoder.utilities import *
+
+import random
+
+
+def lcs(u, v):
+    # t[(n,m)] = length of longest common string ending at first
+    # n elements of u & first m elements of v
+    t = {}
+
+    for n in range(len(u) + 1):
+        for m in range(len(v) + 1):
+            if m == 0 or n == 0:
+                t[(n, m)] = 0
+                continue
+
+            if u[n - 1] == v[m - 1]:
+                t[(n, m)] = 1 + t[(n - 1, m - 1)]
+            else:
+                t[(n, m)] = 0
+    l, n, m = max((l, n, m) for (n, m), l in t.items())
+    return u[n - l:n]
+
+
+delimiters = ['.', ',', ' ', '(', ')', '-']
+characters = [chr(ord('a') + j)
+              for j in range(26)] + \
+             [chr(ord('A') + j)
+              for j in range(26)] + \
+    [str(j) for j in range(10)] + \
+    ['+']
+
+WORDS = None
+
+
+def randomDelimiter():
+    return random.choice(delimiters)
+
+
+def randomCharacter():
+    return random.choice(characters)
+
+
+def randomWord(minimum=1, predicate=None):
+    global WORDS
+    if WORDS is None:
+        tasks, cheating = loadPBETasks()
+        observations = {''.join(z)
+                        for t in tasks
+                        for xs, y in t.examples
+                        for z in list(xs) + [y]}
+
+        def splitMany(s, ds):
+            if len(ds) == 0:
+                return [s]
+            d = ds[0]
+            ds = ds[1:]
+            s = [w
+                 for z in s.split(d)
+                 for w in splitMany(z, ds)
+                 if len(w) > 0]
+            return s
+
+        WORDS = {w
+                 for o in observations
+                 for w in splitMany(o, delimiters)}
+        WORDS = list(sorted(list(WORDS)))
+
+    # a disproportionately large fraction of the words have length three
+    # the purpose of this is to decrease the number of 3-length words we have
+    while True:
+        if random.random() > 0.7:
+            candidate = random.choice([w for w in WORDS if len(w) >= minimum])
+        else:
+            candidate = random.choice(
+                [w for w in WORDS if len(w) >= minimum and len(w) != 3])
+        if predicate is None or predicate(candidate):
+            return candidate
+
+
+def randomWords(ds, minimum=1, lb=2, ub=4):
+    words = [randomWord(minimum=minimum)
+             for _ in range(random.choice(range(lb, ub+1)))]
+    s = ""
+    for j,w in enumerate(words):
+        if j > 0:
+            s += random.choice(ds)
+        s += w
+    return s
+
+
+def makeTasks():
+    import random
+    random.seed(9)
+
+    NUMBEROFEXAMPLES = 4
+
+    problems = []
+
+    def toList(s): return [c for c in s]
+    # Converts strings into a list of characters depending on the type
+
+    def preprocess(x):
+        if isinstance(x, tuple):
+            return tuple(preprocess(z) for z in x)
+        if isinstance(x, list):
+            return [preprocess(z) for z in x]
+        if isinstance(x, str):
+            return [c for c in x]
+        if isinstance(x, bool):
+            return x
+        assert False
+
+    def problem(n, examples, needToTrain=False):
+        task = Task(n, guess_arrow_type(examples),
+                    [(preprocess(x),
+                      preprocess(y))
+                     for x, y in examples])
+        task.mustTrain = True
+        problems.append(task)
+
+    for d1, d2 in randomPermutation(crossProduct(delimiters, delimiters))[
+            :len(delimiters) * 2]:
+        if d1 != d2:
+            problem("Replace '%s' w/ '%s'" % (d1, d2),
+                    [((x,), x.replace(d1, d2))
+                     for _ in range(NUMBEROFEXAMPLES)
+                     for x in [randomWords(d1)]],
+                    needToTrain=False)
+    for d in delimiters:
+        problem("drop first word delimited by '%s'" % d,
+                [((x,), d.join(x.split(d)[1:]))
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for x in [randomWords(d)]],
+                needToTrain=True)
+        for n in [0, 1, -1]:
+            problem("nth (n=%d) word delimited by '%s'" % (n, d),
+                    [((x,), x.split(d)[n])
+                     for _ in range(NUMBEROFEXAMPLES)
+                     for x in [randomWords(d)]],
+                    needToTrain=True)
+    for d1 in delimiters:
+        problem("Append two words delimited by '%s'" % (d1),
+                [((x, y), x + d1 + y)
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for x in [randomWord()]
+                 for y in [randomWord()]],
+                needToTrain=True)
+    for d1, d2 in randomPermutation(
+        crossProduct(
+            delimiters, delimiters))[
+            :len(delimiters)]:
+        problem("Append two words delimited by '%s%s'" % (d1, d2),
+                [((x, y), x + d1 + d2 + y)
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for x in [randomWord()]
+                 for y in [randomWord()]],
+                needToTrain=True)
+    for n in range(1, 6):
+        problem("Drop last %d characters" % n,
+                [((x,), x[:-n])
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for x in [randomWord(minimum=n)]],
+                needToTrain=True)
+        if n > 1:
+            problem("Take first %d characters" % n,
+                    [((x,), x[:n])
+                     for _ in range(NUMBEROFEXAMPLES)
+                     for x in [randomWord(minimum=n)]],
+                    needToTrain=True)
+    for d1, d2 in randomPermutation(
+        crossProduct(
+            delimiters, delimiters))[
+            :len(delimiters)]:
+        problem("Extract word delimited by '%s' - '%s'" % (d1, d2),
+                [((a + d1 + b + d2 + c + d + e,), b)
+                 for _ in range(int(NUMBEROFEXAMPLES / 2))
+                 for d in [d1, d2]
+                 for a in [randomWord()]
+                 for b in [randomWord()]
+                 for c in [randomWord()]
+                 for e in [randomWord()]],
+                needToTrain=True)
+
+    for n in range(len(delimiters)):
+        problem("First letters of words (%s)" % ("I" * (1 + n)),
+                [((x,), "".join(map(lambda z: z[0], x.split(' '))))
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for x in [randomWords(' ')]
+                 ],
+                needToTrain=True)
+        
+    for d in delimiters:
+        problem("Take first character and append '%s'" % d,
+                [((x,), x[0] + d)
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for x in [randomWord()]],
+                needToTrain=True)
+
+    for n in range(len(delimiters)):
+        problem("Abbreviate separate words (%s)" % ("I" * (n + 1)),
+                [((x, y), "%s.%s." % (x[0], y[0]))
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for y in [randomWord()]
+                 for x in [randomWord()]])
+        d = delimiters[n]
+        problem("Abbreviate words separated by '%s'" % d,
+                [((x + d + y,), "%s.%s." % (x[0], y[0]))
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for y in [randomWord()]
+                 for x in [randomWord()]])
+
+    for n in range(len(delimiters)):
+        problem("Append 2 strings (%s)" % ('I' * (n + 1)),
+                [((x, y), x + y)
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for y in [randomWord()]
+                 for x in [randomWord()]],
+                needToTrain=True)
+
+    for n in range(len(delimiters)):
+        w = randomWord(minimum=3)
+        problem("Prepend '%s'" % w,
+                [((x,), w + x)
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for x in [randomWord()]])
+        w = randomWord(minimum=3)
+        problem("Append '%s'" % w,
+                [((x,), x + w)
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for x in [randomWord()]])
+        w = randomWord(minimum=3)
+        problem("Prepend '%s' to first word" % w,
+                [((x + ' ' + y,), w + x)
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for x in [randomWord()]
+                 for y in [randomWord()]])
+
+    for n in range(1,6):
+        problem("parentheses around a single word (%s)"%('I'*n),
+                [((w,),"(%s)"%w)
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for w in [randomWord()] ])
+    problem("parentheses around first word",
+            [((w + " " + s,),"(%s)"%w)
+             for _ in range(NUMBEROFEXAMPLES)
+             for w in [randomWord()]
+             for s in [randomWords(" ")] ])
+    problem("parentheses around second word",
+            [((s,), "(%s)"%(s.split(" ")[1]))
+             for _ in range(NUMBEROFEXAMPLES)
+             for s in [randomWords(" ")] ])
+
+    allowed = [d for d in delimiters if d not in "()"]
+    for d1,d2 in randomPermutation(crossProduct(allowed, allowed))[:len(delimiters)]:
+        problem("parentheses around word delimited by '%s' & '%s'"%(d1,d2),
+                [((prefix + d1 + word + d2 + suffix,),
+                  prefix + d1 + '(' + word + ')' + d2 + suffix)
+                 for _ in range(NUMBEROFEXAMPLES)
+                 for prefix in [randomWords("", lb=0, ub=1)]
+                 for suffix in [randomWords(allowed, ub=2, lb=1)]
+                 for word in [randomWord()] ])
+
+    for n in range(7):
+        w = randomWord(minimum=3)
+        problem("ensure suffix `%s`"%w,
+                [ ((s + (w if f else ""),), s + w)
+                  for _ in range(NUMBEROFEXAMPLES)
+                  for s in [randomWords(" ")]
+                  for f in [random.choice([True,False])] ])
+            
+
+    for p in problems:
+        guessConstantStrings(p)
+
+    return problems
+
+
+def loadPBETasks(directory="PBE_Strings_Track"):
+    """
+    Processes sygus benchmarks into task objects
+    For these benchmarks, all of the constant strings are given to us.
+    In a sense this is cheating
+    Returns (tasksWithoutCheating, tasksWithCheating).
+    NB: Results in paper are done without "cheating"
+    """
+    import os
+    from sexpdata import loads, Symbol
+
+    def findStrings(s):
+        if isinstance(s, list):
+            return [y
+                    for x in s
+                    for y in findStrings(x)]
+        if isinstance(s, str):
+            return [s]
+        return []
+
+    def explode(s):
+        return [c for c in s]
+
+    tasks = []
+    cheatingTasks = []
+    for f in os.listdir(directory):
+        if not f.endswith('.sl'):
+            continue
+        with open(directory + "/" + f, "r") as handle:
+            message = "(%s)" % (handle.read())
+
+        expression = loads(message)
+
+        constants = []
+        name = f
+        examples = []
+        declarative = False
+        for e in expression:
+            if len(e) == 0:
+                continue
+            if e[0] == Symbol('constraint'):
+                e = e[1]
+                assert e[0] == Symbol('=')
+                inputs = e[1]
+                assert inputs[0] == Symbol('f')
+                inputs = inputs[1:]
+                output = e[2]
+                examples.append((inputs, output))
+            elif e[0] == Symbol('synth-fun'):
+                if e[1] == Symbol('f'):
+                    constants += findStrings(e)
+                else:
+                    declarative = True
+                    break
+        if declarative: continue
+        
+        examples = list({(tuple(xs), y) for xs, y in examples})
+
+        task = Task(name, arrow(*[tstr] * (len(examples[0][0]) + 1)),
+                    [(tuple(map(explode, xs)), explode(y))
+                     for xs, y in examples])
+        cheat = task
+
+        tasks.append(task)
+        cheatingTasks.append(cheat)
+
+    for p in tasks:
+        guessConstantStrings(p)
+    return tasks, cheatingTasks
+
+
+def guessConstantStrings(task):
+    if task.request.returns() == tlist(tcharacter):
+        examples = task.examples
+        guesses = {}
+        N = 10
+        T = 2
+        for n in range(min(N, len(examples))):
+            for m in range(n + 1, min(N, len(examples))):
+                y1 = examples[n][1]
+                y2 = examples[m][1]
+                l = ''.join(lcs(y1, y2))
+                if len(l) > 2:
+                    guesses[l] = guesses.get(l, 0) + 1
+
+        task.stringConstants = [g for g, f in guesses.items()
+                                if f >= T]
+    else:
+        task.stringConstants = []
+                    
+
+    task.BIC = 1.
+    task.maxParameters = 1
+
+    task.specialTask = ("stringConstant",
+                        {"maxParameters": task.maxParameters,
+                         "stringConstants": task.stringConstants})
+
+
+if __name__ == "__main__":
+    challenge, _ = loadPBETasks("data/sygus")
+
+    tasks = makeTasks()
+    print(len(tasks), "synthetic tasks")
+    tasks = []
+    for t in tasks + challenge:
+        print(t.name)
+        for xs, y in t.examples:
+            xs = ['"' + "".join(x) + '"' for x in xs]
+            y = "".join(y)
+            print('f(%s) = "%s"' % (", ".join(xs), y))
+        print("\t{%s}" % (t.stringConstants))
+        print()
+    assert False
+    # def maximumLength(x):
+    #     if isinstance(x,list):
+    #         return max([len(x)] + map(maximumLength,x))
+    #     return 1
+
+    # print max(maximumLength(z) for t in tasks
+    #     for (x,),y in t.examples
+    #     for z in [x,y] )
+
+    if len(sys.argv) > 1 and "json" in sys.argv[1]:
+        import json
+        tasks = makeTasks()
+        obj = [t.as_json_dict() for t in tasks]
+        json.dump(obj, sys.stdout)
+    else:
+        as_tex = len(sys.argv) > 1 and "tex" in sys.argv[1]
+        for t in tasks:
+            print(t.name)
+            print(t.request)
+            if as_tex:
+                print("""\\begin{tabular}{ll}
+                \\toprule Input&Output\\\\\\midrule
+        %s
+        \\\\\\bottomrule
+        \\end{tabular}""" % (" \\\\\n ".join(x[0] + " & " + y for x, y in t.examples)))
+            else:
+                for x, y in t.examples:
+                    print(x[0], '\t', y)
+            print()
+        print(len(tasks), "tasks")

dreamcoder/domains/text/textPrimitives.py

@@ -0,0 +1,87 @@
+from dreamcoder.program import *
+from dreamcoder.domains.text.makeTextTasks import delimiters
+
+def _isUpper(x): return x.isupper()
+
+def _increment(x): return x + 1
+
+
+def _decrement(x): return x - 1
+
+
+def _lower(x): return x.lower()
+
+
+def _upper(x): return x.upper()
+
+
+def _capitalize(x): return x.capitalize()
+
+
+def _append(x): return lambda y: x + y
+
+
+def _slice(x): return lambda y: lambda s: s[x:y]
+
+
+def _index(n): return lambda x: x[n]
+
+
+def _map(f): return lambda x: list(map(f, x))
+
+
+def _find(pattern): return lambda s: s.index(pattern)
+
+
+def _replace(original): return lambda replacement: lambda target: target.replace(
+    original, replacement)
+
+
+def _split(delimiter): return lambda s: s.split(delimiter)
+
+
+def _join(delimiter): return lambda ss: delimiter.join(ss)
+
+
+def _identity(x): return x
+#def _reverse(x): return x[::-1]
+
+
+def _strip(x): return x.strip()
+
+
+def _eq(x): return lambda y: x == y
+
+
+specialCharacters = {' ': 'SPACE',
+                     ')': 'RPAREN',
+                     '(': 'LPAREN'}
+
+primitives = [
+    Primitive("char-eq?", arrow(tcharacter, tcharacter, tboolean), _eq),
+    Primitive("STRING", tstr, None)
+] + [Primitive("'%s'" % d, tcharacter, d) for d in delimiters if d not in specialCharacters] + \
+    [Primitive(name, tcharacter, value) for value, name in specialCharacters.items()]
+
+
+def _cons(x): return lambda y: [x] + y
+
+
+def _car(x): return x[0]
+
+
+def _cdr(x): return x[1:]
+
+
+targetTextPrimitives = [
+    Primitive("take-word", arrow(tcharacter, tstr, tstr), None),
+    Primitive("drop-word", arrow(tcharacter, tstr, tstr), None),
+    Primitive("append", arrow(tlist(t0), tlist(t0), tlist(t0)), None),
+    Primitive("abbreviate", arrow(tstr, tstr), None),
+    Primitive("last-word", arrow(tcharacter, tstr, tstr), None),
+    Primitive("replace-character", arrow(tcharacter, tcharacter, tstr, tstr), None),
+] + primitives + [
+    Primitive("empty", tlist(t0), []),
+    Primitive("cons", arrow(t0, tlist(t0), tlist(t0)), _cons),
+    Primitive("car", arrow(tlist(t0), t0), _car),
+    Primitive("cdr", arrow(tlist(t0), tlist(t0)), _cdr)]

dreamcoder/domains/tower/main.py

@@ -0,0 +1,359 @@
+from dreamcoder.dreamcoder import *
+
+from dreamcoder.domains.tower.towerPrimitives import primitives, new_primitives, animateTower
+from dreamcoder.domains.tower.makeTowerTasks import *
+from dreamcoder.domains.tower.tower_common import renderPlan, towerLength, centerTower
+from dreamcoder.utilities import *
+
+import os
+import datetime
+
+try: #pypy will fail
+    from dreamcoder.recognition import variable
+    import torch.nn as nn
+    import torch.nn.functional as F
+
+    class Flatten(nn.Module):
+        def __init__(self):
+            super(Flatten, self).__init__()
+
+        def forward(self, x):
+            return x.view(x.size(0), -1)
+
+
+    class TowerCNN(nn.Module):
+        special = 'tower'
+
+        def __init__(self, tasks, testingTasks=[], cuda=False, H=64):
+            super(TowerCNN, self).__init__()
+            self.CUDA = cuda
+            self.recomputeTasks = True
+
+            self.outputDimensionality = H
+            def conv_block(in_channels, out_channels):
+                return nn.Sequential(
+                    nn.Conv2d(in_channels, out_channels, 3, padding=1),
+                    # nn.BatchNorm2d(out_channels),
+                    nn.ReLU(),
+                    nn.MaxPool2d(2)
+                )
+
+            self.inputImageDimension = 256
+            self.resizedDimension = 64
+            assert self.inputImageDimension % self.resizedDimension == 0
+
+            # channels for hidden
+            hid_dim = 64
+            z_dim = 64
+
+            self.encoder = nn.Sequential(
+                conv_block(6, hid_dim),
+                conv_block(hid_dim, hid_dim),
+                conv_block(hid_dim, hid_dim),
+                conv_block(hid_dim, z_dim),
+                Flatten()
+            )
+
+            self.outputDimensionality = 1024
+
+            if cuda:
+                self.CUDA=True
+                self.cuda()  # I think this should work?
+
+        def forward(self, v, v2=None):
+            """v: tower to build. v2: image of tower we have built so far"""
+            # insert batch if it is not already there
+            if len(v.shape) == 3:
+                v = np.expand_dims(v, 0)
+                inserted_batch = True
+                if v2 is not None:
+                    assert len(v2.shape) == 3
+                    v2 = np.expand_dims(v2, 0)
+            elif len(v.shape) == 4:
+                inserted_batch = False
+                pass
+            else:
+                assert False, "v has the shape %s"%(str(v.shape))
+
+            if v2 is None: v2 = np.zeros(v.shape)
+
+            v = np.concatenate((v,v2), axis=3)
+            v = np.transpose(v,(0,3,1,2))
+            assert v.shape == (v.shape[0], 6,self.inputImageDimension,self.inputImageDimension)
+            v = variable(v, cuda=self.CUDA).float()
+            window = int(self.inputImageDimension/self.resizedDimension)
+            v = F.avg_pool2d(v, (window,window))
+            #showArrayAsImage(np.transpose(v.data.numpy()[0,:3,:,:],[1,2,0]))
+            v = self.encoder(v)
+            if inserted_batch:
+                return v.view(-1)
+            else:
+                return v
+
+        def featuresOfTask(self, t, t2=None):  # Take a task and returns [features]
+            return self(t.getImage(),
+                        None if t2 is None else t2.getImage(drawHand=True))
+
+        def featuresOfTasks(self, ts, t2=None):  # Take a task and returns [features]
+            """Takes the goal first; optionally also takes the current state second"""
+            if t2 is None:
+                pass
+            elif isinstance(t2, Task):
+                assert False
+                #t2 = np.array([t2.getImage(drawHand=True)]*len(ts))
+            elif isinstance(t2, list):
+                t2 = np.array([t.getImage(drawHand=True) if t else np.zeros((self.inputImageDimension,
+                                                                             self.inputImageDimension,
+                                                                             3))
+                               for t in t2])
+            else:
+                assert False
+
+            return self(np.array([t.getImage() for t in ts]),
+                        t2)
+
+        def taskOfProgram(self, p, t,
+                          lenient=False):
+            try:
+                pl = executeTower(p,0.05)
+                if pl is None or (not lenient and len(pl) == 0): return None
+                if len(pl) > 100 or towerLength(pl) > 360: return None
+
+                t = SupervisedTower("tower dream", p)
+                return t
+            except Exception as e:
+                return None
+except: pass
+
+
+
+def tower_options(parser):
+    parser.add_argument("--tasks",
+                        choices=["old","new"],
+                        default="old")
+    parser.add_argument("--visualize",
+                        default=None, type=str)
+    parser.add_argument("--solutions",
+                        default=None, type=str)
+    parser.add_argument("--split",
+                        default=1., type=float)
+    parser.add_argument("--dream",
+                        default=None, type=str)
+    parser.add_argument("--primitives",
+                        default="old", type=str,
+                        choices=["new", "old"])
+
+
+def dreamOfTowers(grammar, prefix, N=250, make_montage=True):
+    request = arrow(ttower,ttower)
+    randomTowers = [tuple(centerTower(t))
+                    for _ in range(N)
+                    for program in [grammar.sample(request,
+                                                   maximumDepth=12,
+                                                   maxAttempts=100)]
+                    if program is not None
+                    for t in [executeTower(program, timeout=0.5) or []]
+                    if len(t) >= 1 and len(t) < 100 and towerLength(t) <= 360.]
+    matrix = [renderPlan(p,Lego=True,pretty=True)
+              for p in randomTowers]
+
+    # Only visualize if it has something to visualize.
+    if len(matrix) > 0:
+        import scipy.misc
+        if make_montage:
+            matrix = montage(matrix)
+            scipy.misc.imsave('%s.png'%prefix, matrix)
+        else:
+            for n,i in enumerate(matrix):
+                scipy.misc.imsave(f'{prefix}/{n}.png', i)
+    else:
+        eprint("Tried to visualize dreams, but none to visualize.")
+
+    
+def visualizePrimitives(primitives, fn=None):
+    from itertools import product
+    #from pylab import imshow,show
+
+    from dreamcoder.domains.tower.towerPrimitives import _left,_right,_loop,_embed,_empty_tower,TowerState
+    _13 = Program.parse("1x3").value
+    _31 = Program.parse("3x1").value
+
+    r = lambda n,k: _right(2*n)(k)
+    l = lambda n,k: _left(2*n)(k)
+    _e = _embed
+    _lp = lambda n,b,k: _loop(n)(b)(k)
+    _arch = lambda k: l(1,_13(r(2,_13(l(1,_31(k))))))
+    _tallArch = lambda h,z,k: _lp(h, lambda _: _13(r(2,_13(l(2,z)))),
+                                  r(1,_31(k)))
+
+    matrix = []
+    for p in primitives:
+        if not p.isInvented: continue
+        eprint(p,":",p.tp)
+        t = p.tp
+        if t.returns() != ttower: continue
+
+        def argumentChoices(t):
+            if t == ttower:
+                return [_empty_tower]
+            elif t == tint:
+                return list(range(5))
+            elif t == arrow(ttower,ttower):
+                return [_arch,_13,_31]
+            else:
+                return []
+
+        ts = []
+        for arguments in product(*[argumentChoices(t) for t in t.functionArguments() ]):
+            t = p.evaluate([])
+            for a in arguments: t = t(a)
+            t = t(TowerState())[1]
+            ts.append(t)
+
+        if ts == []: continue
+        
+        matrix.append([renderPlan(p,pretty=True)
+                       for p in ts])
+
+    # Only visualize if it has something to visualize.
+    if len(matrix) > 0:
+        matrix = montageMatrix(matrix)
+        # imshow(matrix)
+        
+        import scipy.misc
+        scipy.misc.imsave(fn, matrix)
+        #    show()
+    else:
+        eprint("Tried to visualize primitives, but none to visualize.")
+
+def animateSolutions(checkpoint):
+    with open(checkpoint,"rb") as handle: result = dill.load(handle)
+    for n,f in enumerate(result.taskSolutions.values()):
+        animateTower(f"/tmp/tower_animation_{n}",f.bestPosterior.program)
+    
+def visualizeSolutions(solutions, export, tasks=None):
+
+    if tasks is None:
+        tasks = list(solutions.keys())
+        tasks.sort(key=lambda t: len(t.plan))
+
+    matrix = []
+    for t in tasks:
+        i = renderPlan(centerTower(t.plan),pretty=True,Lego=True)
+        if solutions[t].empty: i = i/3.
+        matrix.append(i)
+
+    # Only visualize if it has something to visualize.
+    if len(matrix) > 0:
+        matrix = montage(matrix)
+        import scipy.misc
+        scipy.misc.imsave(export, matrix)
+    else:
+        eprint("Tried to visualize solutions, but none to visualize.")
+
+
+def main(arguments):
+    """
+    Takes the return value of the `commandlineArguments()` function as input and
+    trains/tests the model on a set of tower-building tasks.
+    """
+
+    # The below global statement is required since primitives is modified within main().
+    # TODO(lcary): use a function call to retrieve and declare primitives instead.
+    global primitives
+
+    import scipy.misc
+
+    g0 = Grammar.uniform({"new": new_primitives,
+                          "old": primitives}[arguments.pop("primitives")],
+                         continuationType=ttower)
+
+    checkpoint = arguments.pop("visualize")
+    if checkpoint is not None:
+        with open(checkpoint,'rb') as handle:
+            primitives = pickle.load(handle).grammars[-1].primitives
+        visualizePrimitives(primitives)
+        sys.exit(0)
+    checkpoint = arguments.pop("solutions")
+    if checkpoint is not None:
+        with open(checkpoint,'rb') as handle:
+            solutions = pickle.load(handle).taskSolutions
+        visualizeSolutions(solutions,
+                           checkpoint + ".solutions.png")
+        animateSolutions(checkpoint)
+        sys.exit(0)
+    checkpoint = arguments.pop("dream")
+    if checkpoint is not None:
+        with open(checkpoint,'rb') as handle:
+            g = pickle.load(handle).grammars[-1]
+        os.system("mkdir  -p data/tower_dreams")
+        dreamOfTowers(g, "data/tower_dreams", make_montage=False)
+        sys.exit(0)
+        
+    
+    tasks = arguments.pop("tasks")
+    if tasks == "new":
+        tasks = makeSupervisedTasks()
+    elif tasks == "old":
+        tasks = makeOldSupervisedTasks()
+    else: assert False
+        
+    test, train = testTrainSplit(tasks, arguments.pop("split"))
+    eprint("Split %d/%d test/train" % (len(test), len(train)))
+
+    # Make a montage for the paper
+    shuffledTrain = list(train)
+    shuffledTest = list(test)
+    random.shuffle(shuffledTrain)
+    shuffledTrain = shuffledTrain + [None]*(60 - len(shuffledTrain))
+    random.shuffle(shuffledTest)
+    shuffledTest = shuffledTest + [None]*(60 - len(shuffledTest))
+    try:
+        SupervisedTower.exportMany("/tmp/every_tower.png",shuffledTrain + shuffledTest, shuffle=False, columns=10)
+        for j,task in enumerate(tasks):
+            task.exportImage(f"/tmp/tower_task_{j}.png")
+        for k,v in dSLDemo().items():
+            scipy.misc.imsave(f"/tmp/tower_dsl_{k}.png", v)
+            os.system(f"convert /tmp/tower_dsl_{k}.png -channel RGB -negate /tmp/tower_dsl_{k}.png")
+    except:
+        eprint("WARNING: can't export images. scipy needs to be an older version")
+        
+
+    timestamp = datetime.datetime.now().isoformat()
+    outputDirectory = "experimentOutputs/towers/%s"%timestamp
+    os.system("mkdir -p %s"%outputDirectory)
+
+    os.system("mkdir  -p data/tower_dreams_initial")
+    try:
+        dreamOfTowers(g0, "data/tower_dreams_initial", make_montage=False)
+        dreamOfTowers(g0, "%s/random_0"%outputDirectory)
+    except:
+        eprint("WARNING: can't export images. scipy needs to be an older version")
+
+    evaluationTimeout = 0.005
+    generator = ecIterator(g0, train,
+                           testingTasks=test,
+                           outputPrefix="%s/tower"%outputDirectory,
+                           evaluationTimeout=evaluationTimeout,
+                           **arguments)
+
+    
+    
+    for result in generator:
+        continue
+        iteration = len(result.learningCurve)
+        newTowers = [tuple(centerTower(executeTower(frontier.sample().program)))
+                     for frontier in result.taskSolutions.values() if not frontier.empty]
+        try:
+            fn = '%s/solutions_%d.png'%(outputDirectory,iteration)
+            visualizeSolutions(result.taskSolutions, fn,
+                               train)
+            eprint("Exported solutions to %s\n"%fn)
+            dreamOfTowers(result.grammars[-1],
+                          '%s/random_%d'%(outputDirectory,iteration))           
+        except ImportError:
+            eprint("Could not import required libraries for exporting towers.")
+        primitiveFilename = '%s/primitives_%d.png'%(outputDirectory, iteration)
+        visualizePrimitives(result.grammars[-1].primitives,
+                            primitiveFilename)
+        eprint("Exported primitives to",primitiveFilename)

dreamcoder/domains/tower/makeTowerTasks.py

@@ -0,0 +1,556 @@
+from dreamcoder.domains.tower.towerPrimitives import ttower, executeTower, _empty_tower, TowerState
+from dreamcoder.domains.tower.tower_common import renderPlan
+from dreamcoder.task import *
+
+
+class SupervisedTower(Task):
+    def __init__(self, name, program, mustTrain=False):
+        if isinstance(program,str):
+            try:
+                program = parseTower(program)
+            except:
+                eprint("Parse failure:")
+                eprint(program)
+                assert False
+            self.original = program
+            plan = executeTower(program)
+        elif isinstance(program,Program):
+            self.original = program
+            plan = executeTower(program)
+        else:
+            plan = program
+        self.original = program
+        state, self.plan = program.evaluate([])(_empty_tower)(TowerState())
+        self.hand = state.hand
+        super(SupervisedTower, self).__init__(name, arrow(ttower,ttower), [],
+                                              features=[])
+        self.specialTask = ("supervisedTower",
+                            {"plan": self.plan})
+        self.image = None
+        self.handImage = None
+        self.mustTrain = mustTrain
+
+    def getImage(self, drawHand=False, pretty=False):
+        if not drawHand:
+            if not pretty:
+                if self.image is not None: return self.image
+                self.image = renderPlan(self.plan, pretty=pretty)
+                return self.image
+            else:
+                return renderPlan(self.plan, pretty=True)
+        else:
+            if self.handImage is not None: return self.handImage
+            self.handImage = renderPlan(self.plan,
+                                        drawHand=self.hand,
+                                        pretty=pretty)
+            return self.handImage
+                
+
+    
+    # do not pickle the image
+    def __getstate__(self):
+        return self.specialTask, self.plan, self.request, self.cache, self.name, self.examples
+    def __setstate__(self, state):
+        self.specialTask, self.plan, self.request, self.cache, self.name, self.examples = state
+        self.image = None
+
+
+    def animate(self):
+        from pylab import imshow,show
+        a = renderPlan(self.plan)
+        imshow(a)
+        show()
+
+    @staticmethod
+    def showMany(ts):
+        from pylab import imshow,show
+        a = montage([renderPlan(t.plan, pretty=True, Lego=True, resolution=256,
+                                drawHand=False)
+                     for t in ts]) 
+        imshow(a)
+        show()
+
+    @staticmethod
+    def exportMany(f, ts, shuffle=True, columns=None):
+        import numpy as np
+        
+        ts = list(ts)
+        if shuffle:
+            assert all( t is not None for t in ts  )
+            random.shuffle(ts)
+        a = montage([renderPlan(t.plan, pretty=True, Lego=True, resolution=256) if t is not None \
+                     else np.zeros((256,256,3))
+                     for t in ts],
+                    columns=columns)        
+        import scipy.misc
+        scipy.misc.imsave(f, a)
+        
+
+    def exportImage(self, f, pretty=True, Lego=True, drawHand=False):
+        a = renderPlan(self.plan,
+                       pretty=pretty, Lego=Lego,
+                       drawHand=t.hand if drawHand else None)
+        import scipy.misc
+        scipy.misc.imsave(f, a)
+
+    def logLikelihood(self, e, timeout=None):
+        from dreamcoder.domains.tower.tower_common import centerTower
+        yh = executeTower(e, timeout)
+        if yh is not None and centerTower(yh) == centerTower(self.plan): return 0.
+        return NEGATIVEINFINITY
+
+    
+
+def parseTower(s):
+    _13 = Program.parse("1x3")
+    _31 = Program.parse("3x1")
+    _r = Program.parse("right")
+    _l = Program.parse("left")
+    _addition = Program.parse("+")
+    _subtraction = Program.parse("-")
+    _lp = Program.parse("tower_loopM")
+    _e = Program.parse("tower_embed")
+
+    from sexpdata import loads, Symbol
+    s = loads(s)
+    def command(k, environment, continuation):
+        if k == Symbol("1x3") or k == Symbol("v"): return Application(_13, continuation)
+        if k == Symbol("3x1") or k == Symbol("h"): return Application(_31, continuation)
+        assert isinstance(k,list)
+        if k[0] == Symbol("r"): return Application(Application(_r, expression(k[1],environment)),continuation)
+        if k[0] == Symbol("l"): return Application(Application(_l, expression(k[1],environment)),continuation)
+        if k[0] == Symbol("for"):
+            v = k[1]
+            b = expression(k[2], environment)
+            newEnvironment = [None, v] + environment
+            body = block(k[3:], newEnvironment, Index(0))
+            return Application(Application(Application(_lp,b),
+                                           Abstraction(Abstraction(body))),
+                               continuation)
+        if k[0] == Symbol("embed"):
+            body = block(k[1:], [None] + environment, Index(0))
+            return Application(Application(_e,Abstraction(body)),continuation)
+            
+        assert False
+    def expression(e, environment):
+        for n, v in enumerate(environment):
+            if e == v: return Index(n)
+
+        if isinstance(e,int): return Program.parse(str(e))
+
+        assert isinstance(e,list)
+        if e[0] == Symbol('+'): return Application(Application(_addition, expression(e[1], environment)),
+                                                   expression(e[2], environment))
+        if e[0] == Symbol('-'): return Application(Application(_subtraction, expression(e[1], environment)),
+                                                   expression(e[2], environment))
+        assert False
+        
+    def block(b, environment, continuation):
+        if len(b) == 0: return continuation
+        return command(b[0], environment, block(b[1:], environment, continuation))
+
+    try: return Abstraction(command(s, [], Index(0)))
+    except: return Abstraction(block(s, [], Index(0)))
+
+    
+def makeSupervisedTasks():
+    arches = [SupervisedTower("arch leg %d"%n,
+                              "((for i %d v) (r 4) (for i %d v) (l 2) h)"%(n,n))
+              for n in range(1,9)
+    ]
+    archesStacks = [SupervisedTower("arch stack %d"%n,
+                                    """
+                                    (for i %d 
+                                    v (r 4) v (l 2) h (l 2))
+                                    """%n)
+                    for n in range(3,7) ]
+    Bridges = [SupervisedTower("bridge (%d) of arch %d"%(n,l),
+                               """
+                               (for j %d
+                                (for i %d 
+                                 v (r 4) v (l 4)) (r 2) h 
+                                (r 4))
+                               """%(n,l))
+               for n in range(2,8)
+               for l in range(1,6)]
+    offsetArches = [SupervisedTower("bridge (%d) of arch, spaced %d"%(n,l),
+                                    """
+                                    (for j %d
+                                    (embed v (r 4) v (l 2) h )
+                                    (r %d))
+                                    """%(n,l),
+                                    mustTrain=n == 3)
+                    for n,l in [(3,7),(4,8)]]
+    Josh = [SupervisedTower("Josh (%d)"%n,
+                            """(for i %d
+                            h (l 2) v (r 2) v (r 2) v (l 2) h (r 6))"""%n)
+            for n in range(1,7) ]
+    
+    staircase1 = [SupervisedTower("R staircase %d"%n,
+"""
+(for i %d (for j i
+(embed v (r 4) v (l 2) h)) (r 6))
+"""%(n))
+                 for n in range(3,8) ]
+    staircase2 = [SupervisedTower("L staircase %d"%n,
+"""
+(for i %d (for j i
+(embed v (r 4) v (l 2) h)) (l 6))
+"""%(n))
+                 for n in range(3,8) ]
+    simpleLoops = [SupervisedTower("%s row %d, spacing %d"%(o,n,s),
+                                   """(for j %d %s (r %s))"""%(n,o,s),
+                                   mustTrain=True)
+                   for o,n,s in [('h',4,7), ('v',5,3)] ]
+
+    pyramids = []
+    pyramids += [SupervisedTower("arch pyramid %d"%n,
+                                 """((for i %d (for j i (embed v (r 4) v (l 2) h)) (r 6))
+                                 (for i %d (for j (- %d i) (embed v (r 4) v (l 2) h)) (r 6)))"""%(n,n,n))
+                for n in range(2,6) ]
+    pyramids += [SupervisedTower("H pyramid %d"%n,
+                                 """((for i %d (for j i h) (r 6))
+                                 (for i %d (for j (- %d i) h) (r 6)))"""%(n,n,n))
+                for n in range(4,6) ]
+#     pyramids += [SupervisedTower("V pyramid %d"%n,
+# """
+# ((for i %d (for j i v) (r 2))
+#  (for i %d (for j (- %d i) v) (r 2)))
+# """%(n,n,n))
+#                 for n in range(4,8) ]
+#     pyramids += [SupervisedTower("V3 pyramid %d"%n,
+# """
+# ((for i %d (for j i v) (r 6))
+#  (for i %d (for j (- %d i) v) (r 6)))
+# """%(n,n,n))
+#                  for n in range(4,8) ]
+    pyramids += [SupervisedTower("H 1/2 pyramid %d"%n,
+                                 """
+(for i %d
+  (r 6)
+  (embed
+    (for j i h (l 3))))
+                                 """%n)
+                for n in range(4,8) ]
+    pyramids += [SupervisedTower("arch 1/2 pyramid %d"%n,
+"""
+(for i %d
+  (r 6)
+  (embed
+    (for j i (embed v (r 4) v (l 2) h) (l 3))))
+"""%n)
+                for n in range(2,8) ]
+    if False:
+        pyramids += [SupervisedTower("V 1/2 pyramid %d"%n,
+                                     """
+    (for i %d
+      (r 2)
+      (embed
+                                     (for j i v (l 1))))"""%(n))
+                    for n in range(4,8) ]
+    bricks = [SupervisedTower("brickwall, %dx%d"%(w,h),
+                              """(for j %d
+                              (embed (for i %d h (r 6)))
+                              (embed (r 3) (for i %d h (r 6))))"""%(h,w,w))
+              for w in range(3,7)
+              for h in range(1,6) ]
+    aqueducts = [SupervisedTower("aqueduct: %dx%d"%(w,h),
+                                 """(for j %d
+                                 %s (r 4) %s (l 2) h (l 2) v (r 4) v (l 2) h (r 4))"""%
+                                 (w, "v "*h, "v "*h))
+                 for w in range(4,8)
+                 for h in range(3,6)
+                 ]
+
+    compositions = [SupervisedTower("%dx%d-bridge on top of %dx%d bricks"%(b1,b2,w1,w2),
+                                    """
+                                    ((for j %d
+                                    (embed (for i %d h (r 6)))
+                                    (embed (r 3) (for i %d h (r 6))))
+                                    (r 1)
+                                    (for j %d
+                                    (for i %d 
+                                    v (r 4) v (l 4)) (r 2) h 
+                                    (r 4)))
+                                    """%(w1,w2,w2,b1,b2))
+                    for b1,b2,w1,w2 in [(5,2,4,5)]
+                    ] + [
+                        SupervisedTower("%d pyramid on top of %dx%d bricks"%(p,w1,w2),
+                                        """
+                                        ((for j %d
+                                        (embed (for i %d h (r 6)))
+                                        (embed (r 3) (for i %d h (r 6))))
+                                        (r 1)
+                                        (for i %d (for j i (embed v (r 4) v (l 2) h)) (r 6))
+                                        (for i %d (for j (- %d i) (embed v (r 4) v (l 2) h)) (r 6)))
+                                        """%(w1,w2,w2,p,p,p))
+                        for w1,w2,p in [(2,5,2)]
+                        ] + \
+                        [
+                            SupervisedTower("%d tower on top of %dx%d bricks"%(t,w1,w2),
+                                            """
+                                            ((for j %d
+                                            (embed (for i %d h (r 6)))
+                                            (embed (r 3) (for i %d h (r 6))))
+                                            (r 6)
+                                            %s (r 4) %s (l 2) h)
+                                            """%(w1,w2,w2,
+                                                 "v "*t, "v "*t))
+                            for t,w1,w2 in [(4,1,3)] ]
+                            
+    
+                     
+    everything = arches + simpleLoops + Bridges + archesStacks + aqueducts + offsetArches + pyramids + bricks + staircase2 + staircase1 + compositions
+    if False:
+        for t in everything:
+            delattr(t,'original')
+    return everything
+
+def makeOldSupervisedTasks():
+    arches = [SupervisedTower("arch leg %d"%n,
+                              "((for i %d v) (r 4) (for i %d v) (l 2) h)"%(n,n))
+              for n in range(1,9)
+    ]
+    archesStacks = [SupervisedTower("arch stack %d"%n,
+                                    """
+                                    (for i %d 
+                                    v (r 4) v (l 2) h (l 2))
+                                    """%n)
+                    for n in range(3,7) ]
+    Bridges = [SupervisedTower("bridge (%d) of arch %d"%(n,l),
+                               """
+                               (for j %d
+                                (for i %d 
+                                 v (r 4) v (l 4)) (r 2) h 
+                                (r 4))
+                               """%(n,l))
+               for n in range(2,8)
+               for l in range(1,6)]
+    offsetArches = [SupervisedTower("bridge (%d) of arch, spaced %d"%(n,l),
+                               """
+                               (for j %d
+                                 v (r 4) v (l 2) h 
+                                (r %d))
+                               """%(n,l))
+                    for n,l in [(3,7),(4,6)]]
+    Josh = [SupervisedTower("Josh (%d)"%n,
+                            """(for i %d
+                            h (l 2) v (r 2) v (r 2) v (l 2) h (r 6))"""%n)
+            for n in range(1,7) ]
+    
+    staircase1 = [SupervisedTower("R staircase %d"%n,
+"""
+(for i %d (for j i
+(embed v (r 4) v (l 2) h)) (r 6))
+"""%(n))
+                 for n in range(3,8) ]
+    staircase2 = [SupervisedTower("L staircase %d"%n,
+"""
+(for i %d (for j i
+(embed v (r 4) v (l 2) h)) (l 6))
+"""%(n))
+                 for n in range(3,8) ]
+    simpleLoops = [SupervisedTower("horizontal row %d, spacing %d"%(n,s),
+                                   """(for j %d h (r %s))"""%(n,s))
+                   for n,s in [(4,6),(5,7)] ]+\
+                [SupervisedTower("horizontal stack %d"%n,
+                                   """(for j %d h)"""%n)
+                   for n in range(5,8) ]+\
+                [SupervisedTower("vertical stack %d"%n,
+                                   """(for j %d v)"""%n)
+                   for n in [5,7] ]
+    pyramids = []
+    pyramids += [SupervisedTower("arch pyramid %d"%n,
+                                 """((for i %d (for j i (embed v (r 4) v (l 2) h)) (r 6))
+                                 (for i %d (for j (- %d i) (embed v (r 4) v (l 2) h)) (r 6)))"""%(n,n,n))
+                for n in range(2,6) ]
+    pyramids += [SupervisedTower("H pyramid %d"%n,
+                                 """((for i %d (for j i h) (r 6))
+                                 (for i %d (for j (- %d i) h) (r 6)))"""%(n,n,n))
+                for n in range(4,6) ]
+#     pyramids += [SupervisedTower("V pyramid %d"%n,
+# """
+# ((for i %d (for j i v) (r 2))
+#  (for i %d (for j (- %d i) v) (r 2)))
+# """%(n,n,n))
+#                 for n in range(4,8) ]
+#     pyramids += [SupervisedTower("V3 pyramid %d"%n,
+# """
+# ((for i %d (for j i v) (r 6))
+#  (for i %d (for j (- %d i) v) (r 6)))
+# """%(n,n,n))
+#                  for n in range(4,8) ]
+    pyramids += [SupervisedTower("H 1/2 pyramid %d"%n,
+                                 """
+(for i %d
+  (r 6)
+  (embed
+    (for j i h (l 3))))
+                                 """%n)
+                for n in range(4,8) ]
+    pyramids += [SupervisedTower("arch 1/2 pyramid %d"%n,
+"""
+(for i %d
+  (r 6)
+  (embed
+    (for j i (embed v (r 4) v (l 2) h) (l 3))))
+"""%n)
+                for n in range(2,8) ]
+    if False:
+        pyramids += [SupervisedTower("V 1/2 pyramid %d"%n,
+                                     """
+    (for i %d
+      (r 2)
+      (embed
+                                     (for j i v (l 1))))"""%(n))
+                    for n in range(4,8) ]
+    bricks = [SupervisedTower("brickwall, %dx%d"%(w,h),
+                              """(for j %d
+                              (embed (for i %d h (r 6)))
+                              (embed (r 3) (for i %d h (r 6))))"""%(h,w,w))
+              for w in range(3,7)
+              for h in range(1,6) ]
+    aqueducts = [SupervisedTower("aqueduct: %dx%d"%(w,h),
+                                 """(for j %d
+                                 %s (r 4) %s (l 2) h (l 2) v (r 4) v (l 2) h (r 4))"""%
+                                 (w, "v "*h, "v "*h))
+                 for w in range(4,8)
+                 for h in range(3,6)
+                 ]
+
+    compositions = [SupervisedTower("%dx%d-bridge on top of %dx%d bricks"%(b1,b2,w1,w2),
+                                    """
+                                    ((for j %d
+                                    (embed (for i %d h (r 6)))
+                                    (embed (r 3) (for i %d h (r 6))))
+                                    (r 1)
+                                    (for j %d
+                                    (for i %d 
+                                    v (r 4) v (l 4)) (r 2) h 
+                                    (r 4)))
+                                    """%(w1,w2,w2,b1,b2))
+                    for b1,b2,w1,w2 in [(5,2,4,5)]
+                    ] + [
+                        SupervisedTower("%d pyramid on top of %dx%d bricks"%(p,w1,w2),
+                                        """
+                                        ((for j %d
+                                        (embed (for i %d h (r 6)))
+                                        (embed (r 3) (for i %d h (r 6))))
+                                        (r 1)
+                                        (for i %d (for j i (embed v (r 4) v (l 2) h)) (r 6))
+                                        (for i %d (for j (- %d i) (embed v (r 4) v (l 2) h)) (r 6)))
+                                        """%(w1,w2,w2,p,p,p))
+                        for w1,w2,p in [(2,5,2)]
+                        ] + \
+                        [
+                            SupervisedTower("%d tower on top of %dx%d bricks"%(t,w1,w2),
+                                            """
+                                            ((for j %d
+                                            (embed (for i %d h (r 6)))
+                                            (embed (r 3) (for i %d h (r 6))))
+                                            (r 6)
+                                            %s (r 4) %s (l 2) h)
+                                            """%(w1,w2,w2,
+                                                 "v "*t, "v "*t))
+                            for t,w1,w2 in [(4,1,3)] ]
+                            
+    
+                     
+    everything = arches + simpleLoops + Bridges + archesStacks + aqueducts + offsetArches + pyramids + bricks + staircase2 + staircase1 + compositions
+    if False:
+        for t in everything:
+            delattr(t,'original')
+    return everything
+
+def dSLDemo():
+    DSL = {}
+    bricks = Program.parse("(lambda (lambda (tower_loopM $0 (lambda (lambda (moveHand 3 (reverseHand (tower_loopM $3 (lambda (lambda (moveHand 6 (3x1 $0)))) $0))))))))")
+    DSL["bricks"] = [ [bricks.runWithArguments([x,y + 4,_empty_tower,TowerState()])[1]
+                       for y in range(6, 6 + 3*4, 3) ]
+                      for x in [3,8] ]
+    dimensionality = {}
+    dimensionality["bricks"] = 2
+
+    bridge = Program.parse("(lambda (lambda (tower_loopM $0 (lambda (lambda (#(lambda (#(lambda (lambda (lambda (tower_loopM $0 (lambda (lambda (1x3 (moveHand 4 ($3 $0))))) (moveHand 2 (3x1 $2)))))) $0 (lambda (reverseHand $0)))) (moveHand 4 $0) $3))))))")
+    DSL["bridge"] = [ [bridge.runWithArguments([x,y,_empty_tower,TowerState()])[1]
+                       for x in range(4,4 + 2*4,2) ]
+                       for y in [4,9] ]
+    dimensionality["bridge"] = 2
+
+    staircase = Program.parse("(lambda (tower_loopM $0 (lambda (lambda (#(lambda (lambda (tower_loopM $1 (lambda (lambda (tower_embed (lambda (#(lambda (1x3 (moveHand 4 (1x3 (reverseHand (moveHand 2 (3x1 $0))))))) $0)) $0))) $0))) $1 (moveHand 6 $0))))))")
+    DSL["staircase"] = [ staircase.runWithArguments([n,_empty_tower,TowerState()])[1]
+                         for n in range(4,5 + 3) ]
+
+    pyramid = Program.parse("(lambda (tower_loopM $0 (lambda (lambda (moveHand 6 (tower_embed (lambda (reverseHand ((lambda (lambda (tower_loopM $1 (lambda (lambda (moveHand $2 (1x3 (moveHand 2 (tower_embed (lambda (moveHand 2 (1x3 $0))) (3x1 $0)))))))))) $2 1 $0))) $0))))))")
+    DSL["pyramid"] = [ pyramid.runWithArguments([n,_empty_tower,TowerState()])[1]
+                       for n in range(4,5 + 3) ]
+
+    towerArch = Program.parse("(lambda (lambda ((lambda ((lambda (lambda (lambda (tower_loopM $0 (lambda (lambda (1x3 (moveHand 4 ($3 $0))))) (moveHand 2 (3x1 $2)))))) $0 (lambda (reverseHand (1x3 $0))))) $0 $1)))")
+    DSL["towerArch"] = [ towerArch.runWithArguments([n,_empty_tower,TowerState()])[1]
+                       for n in range(4,5 + 3) ]
+
+    images = {}
+    for k,v in DSL.items():
+        d = dimensionality.get(k,1)
+        if d == 1:
+            i = montageMatrix([[renderPlan(p, pretty=True, Lego=True) for p in v]])
+        elif d == 2:
+            i = montageMatrix([[renderPlan(p, pretty=True, Lego=True) for p in ps] for ps in v] )
+        else: assert False
+
+        images[k] = i
+
+    return images
+            
+if __name__ == "__main__":
+    from pylab import imshow,show
+    from dreamcoder.domains.tower.tower_common import *
+    
+    ts = makeSupervisedTasks()
+    print(len(ts),"total tasks")
+    print("maximum plan length",max(len(f.plan) for f in ts ))
+    print("maximum tower length",max(towerLength(f.plan) for f in ts ))
+    print("maximum tower height",max(towerHeight(simulateWithoutPhysics(f.plan)) for f in ts ))
+    SupervisedTower.exportMany("/tmp/every_tower.png",ts,shuffle=False)
+    
+    for j,t in enumerate(ts):
+        t.exportImage("/tmp/tower_%d.png"%j,
+                      drawHand=False)
+
+    for k,v in dSLDemo().items():
+        import scipy.misc
+        scipy.misc.imsave(f"/tmp/tower_dsl_{k}.png", v)
+
+    exampleTowers = [103,104,105,93,73,
+                     50,67,35,43,106]
+    SupervisedTower.exportMany("/tmp/tower_montage.png",
+                               [ts[n] for n in exampleTowers ],
+                               columns=5,
+                               shuffle=False)
+    assert False
+
+
+    keywords = ["pyramid",
+                "on top of",
+                "arch 1/2 pyramid",
+                "brickwall",
+                "staircase",
+                "bridge",
+                "aqueduct",
+                "spaced",
+                "spaced",
+                "arch stack"]
+    for n in range(100):
+        examples = []
+        for kw in keywords:
+            if kw == "on top of":
+                examples = examples + list(filter(lambda t: kw in str(t), ts))
+            else:
+                examples.append(random.choice(list(filter(lambda t: kw in str(t), ts))))
+
+        random.shuffle(examples)
+        SupervisedTower.exportMany("/tmp/tower10_%d.png"%n,examples,
+                                   columns=int(len(examples)/2))
+        
+        
+        

dreamcoder/domains/tower/towerPrimitives.py

@@ -0,0 +1,152 @@
+from dreamcoder.program import *
+
+
+class TowerState:
+    def __init__(self, hand=0, orientation=1, history=None):
+        # List of (State|Block)
+        self.history = history
+        self.hand = hand
+        self.orientation = orientation
+    def __str__(self): return f"S(h={self.hand},o={self.orientation})"
+    def __repr__(self): return str(self)
+    def left(self, n):
+        return TowerState(hand=self.hand - n, orientation=self.orientation,
+                          history=self.history if self.history is None \
+                          else self.history + [self])
+    def right(self, n): return TowerState(hand=self.hand + n, orientation=self.orientation,
+                                          history=self.history if self.history is None \
+                                          else self.history + [self])
+    def reverse(self): return TowerState(hand=self.hand, orientation=-1*self.orientation,
+                                         history=self.history if self.history is None \
+                                         else self.history + [self])
+    def move(self, n): return TowerState(hand=self.hand + n*self.orientation, orientation=self.orientation,
+                                         history=self.history if self.history is None \
+                                         else self.history + [self])
+
+    def recordBlock(self, b):
+        if self.history is None: return self
+        return TowerState(hand=self.hand,
+                          orientation=self.orientation,
+                          history=self.history + [b])
+        
+
+def _empty_tower(h): return (h,[])
+def _left(d):
+    return lambda k: lambda s: k(s.left(d))
+def _right(d):
+    return lambda k: lambda s: k(s.right(d))
+def _loop(n):
+    def f(start, stop, body, state):
+        if start >= stop: return state,[]
+        state, thisIteration = body(start)(state)
+        state, laterIterations = f(start + 1, stop, body, state)
+        return state, thisIteration + laterIterations
+    def sequence(b,k,h):
+        h,bodyBlocks = f(0,n,b,h)
+        h,laterBlocks = k(h)
+        return h,bodyBlocks+laterBlocks
+    return lambda b: lambda k: lambda h: sequence(b,k,h)
+def _simpleLoop(n):
+    def f(start, body, k):
+        if start >= n: return k
+        return body(start)(f(start + 1, body, k))
+    return lambda b: lambda k: f(0,b,k)
+def _embed(body):
+    def f(k):
+        def g(hand):
+            bodyHand, bodyActions = body(_empty_tower)(hand)
+            # Record history if we are doing that
+            if hand.history is not None:
+                hand = TowerState(hand=hand.hand,
+                                  orientation=hand.orientation,
+                                  history=bodyHand.history)
+            hand, laterActions = k(hand)
+            return hand, bodyActions + laterActions
+        return g
+    return f
+def _moveHand(n):
+    return lambda k: lambda s: k(s.move(n))
+def _reverseHand(k):
+    return lambda s: k(s.reverse())
+    
+class TowerContinuation(object):
+    def __init__(self, x, w, h):
+        self.x = x
+        self.w = w*2
+        self.h = h*2
+    def __call__(self, k):
+        def f(hand):
+            thisAction = [(self.x + hand.hand,self.w,self.h)]
+            hand = hand.recordBlock(thisAction[0])
+            hand, rest = k(hand)
+            return hand, thisAction + rest
+        return f
+
+# name, dimensions
+blocks = {
+    # "1x1": (1.,1.),
+    # "2x1": (2.,1.),
+    # "1x2": (1.,2.),
+    "3x1": (3, 1),
+    "1x3": (1, 3),
+    #          "4x1": (4.,1.),
+    #          "1x4": (1.,4.)
+}
+
+
+ttower = baseType("tower")
+common_primitives = [
+    Primitive("tower_loopM", arrow(tint, arrow(tint, ttower, ttower), ttower, ttower), _simpleLoop),
+    Primitive("tower_embed", arrow(arrow(ttower,ttower), ttower, ttower), _embed),
+] + [Primitive(name, arrow(ttower,ttower), TowerContinuation(0, w, h))
+     for name, (w, h) in blocks.items()] + \
+         [Primitive(str(j), tint, j) for j in range(1,9) ]
+primitives = common_primitives + [
+    Primitive("left", arrow(tint, ttower, ttower), _left),
+    Primitive("right", arrow(tint, ttower, ttower), _right)
+    ]
+
+new_primitives = common_primitives + [
+    Primitive("moveHand", arrow(tint, ttower, ttower), _moveHand),
+    Primitive("reverseHand", arrow(ttower, ttower), _reverseHand)
+    ]
+
+def executeTower(p, timeout=None):
+    try:
+        return runWithTimeout(lambda : p.evaluate([])(_empty_tower)(TowerState())[1],
+                              timeout=timeout)
+    except RunWithTimeout: return None
+    except: return None
+
+def animateTower(exportPrefix, p):
+    print(exportPrefix, p)
+    from dreamcoder.domains.tower.tower_common import renderPlan
+    state,actions = p.evaluate([])(_empty_tower)(TowerState(history=[]))
+    print(actions)
+    trajectory = state.history + [state]
+    print(trajectory)
+    print()
+
+    assert tuple(z for z in trajectory if not isinstance(z, TowerState) ) == tuple(actions)        
+
+    def hd(n):
+        h = 0
+        for state in trajectory[:n]:
+            if isinstance(state, TowerState):
+                h = state.hand
+        return h
+    animation = [renderPlan([b for b in trajectory[:n] if not isinstance(b, TowerState)],
+                            pretty=True, Lego=True,
+                            drawHand=hd(n),
+                            masterPlan=actions,
+                            randomSeed=hash(exportPrefix))
+                 for n in range(0,len(trajectory) + 1)]
+    import scipy.misc
+    import random
+    r = random.random()
+    paths = []
+    for n in range(len(animation)):
+        paths.append(f"{exportPrefix}_{n}.png")
+        scipy.misc.imsave(paths[-1], animation[n])
+    os.system(f"convert -delay 10 -loop 0 {' '.join(paths)} {exportPrefix}.gif")
+#    os.system(f"rm {' '.join(paths)}")

dreamcoder/domains/tower/tower_common.py

@@ -0,0 +1,173 @@
+import random
+import math
+from dreamcoder.utilities import *
+
+def simulateWithoutPhysics(plan,ordered=True):
+    def overlap(b1,
+                b2):
+        (x,w,h) = b1
+        (x_,y_,w_,h_) = b2
+        x1 = x - w/2
+        x2 = x + w/2
+        x1_ = x_ - w_/2
+        x2_ = x_ + w_/2
+        if x1_ >= x2 or x1 >= x2_: return None
+        assert h%2 == 0 and h_%2 == 0
+        return y_ + h_//2 + h//2
+    def lowestPossibleHeight(b):
+        h = b[2]
+        assert h%2 == 0
+        return int(h/2)
+    def placeAtHeight(b,y):
+        (x,w,h) = b
+        return (x,y,w,h)
+    def placeBlock(world, block):
+        lowest = max([lowestPossibleHeight(block)] + \
+                     [overlap(block,other)
+                      for other in world
+                      if overlap(block,other) is not None])
+        world.append(placeAtHeight(block, lowest))
+
+    w = []
+    for p in plan: placeBlock(w,p)
+    if ordered: w = list(sorted(w))
+    return w
+
+def centerTower(t,hand=None, masterPlan=None):
+
+    if len(t) == 0:
+        if hand is None:
+            return t
+        else:
+            return t, hand
+    def getCenter(t):
+        x1 = max(x for x, _, _ in t)
+        x0 = min(x for x, _, _ in t)
+        c = int((x1 - x0) / 2.0) + x0
+        return c
+    c = getCenter(masterPlan or t)
+    t = [(x - c, w, h) for x, w, h in t]
+    if hand is None:
+        return t
+    else:
+        return t, hand - c
+
+def towerLength(t):
+    if len(t) == 0: return 0
+    x1 = max(x for x, _, _ in t)
+    x0 = min(x for x, _, _ in t)
+    return x1 - x0
+
+def towerHeight(t):
+    y1 = max(y + h/2 for _, y, _, h in t )
+    y0 = min(y - h/2 for _, y, _, h in t )
+    return y1 - y0
+
+
+
+def renderPlan(plan, resolution=256, window=64, floorHeight=2, borderSize=1, bodyColor=(0.,1.,1.),
+               borderColor=(1.,0.,0.),
+               truncate=None, randomSeed=None,
+               masterPlan=None,
+               pretty=False, Lego=False,
+               drawHand=None):
+    import numpy as np
+
+    if Lego: assert pretty
+
+    if drawHand is not None and drawHand is not False:
+        plan, drawHand = centerTower(plan, drawHand,
+                                     masterPlan=masterPlan)
+    else:
+        plan = centerTower(plan,masterPlan=masterPlan)
+
+    world = simulateWithoutPhysics(plan,
+                                   ordered=randomSeed is None)
+    if truncate is not None: world = world[:truncate]
+    a = np.zeros((resolution, resolution, 3))
+
+    def transform(x,y):
+        y = resolution - y*resolution/float(window)
+        x = resolution/2 + x*resolution/float(window)
+        return int(x + 0.5),int(y + 0.5)
+    def clip(p):
+        if p < 0: return 0
+        if p >= resolution: return resolution - 1
+        return int(p + 0.5)
+    def clear(x,y):
+        for xp,yp,wp,hp in world:
+            if x < xp + wp/2. and \
+               x > xp - wp/2. and \
+               y < yp + hp/2. and \
+               y > yp - hp/2.:
+                return False
+        return True
+    def bump(x,y,c):
+        size = 0.5*resolution/window
+        x,y = transform(x,y)
+        y -= floorHeight
+        y1 = y
+        y2 = y - size
+        x1 = x - size/2
+        x2 = x + size/2
+        a[clip(y2) : clip(y1),
+          clip(x1) : clip(x2),
+          :] = c
+
+
+    if randomSeed is not None:
+        randomNumbers = random.Random(randomSeed)
+    def _color():
+        if randomSeed is None:
+            return random.random()*0.7 + 0.3
+        else:
+            return randomNumbers.random()*0.7 + 0.3
+    def color():
+        return (_color(),_color(),_color())
+    
+    def rectangle(x1,x2,y1,y2,c,cp=None):
+        x1,y1 = transform(x1,y1)
+        x2,y2 = transform(x2,y2)
+        y1 -= floorHeight
+        y2 -= floorHeight
+        a[clip(y2) : clip(y1),
+          clip(x1) : clip(x2),
+          :] = c
+        if cp is not None:
+            a[clip(y2 + borderSize) : clip(y1 - borderSize),
+              clip(x1 + borderSize) : clip(x2 - borderSize),
+              :] = cp
+
+    for x,y,w,h in world:
+        x1,y1 = x - w/2., y - h/2.                          
+        x2,y2 = x + w/2., y + h/2.
+        if pretty:
+            thisColor = color()
+            rectangle(x1,x2,y1,y2,
+                      thisColor)
+            if Lego:
+                bumps = w
+                for nb in range(bumps):
+                    nx = x - w/2. + 0.5 + nb
+                    ny = y + h/2. + 0.00001
+                    if clear(nx,ny):
+                        bump(nx,ny,thisColor)
+        else:
+            rectangle(x1,x2,y1,y2,
+                      borderColor, bodyColor)
+    
+    a[resolution - floorHeight:,:,:] = 1.
+    if drawHand is not None:
+        if not Lego:
+            dh = 0.25
+            rectangle(drawHand - dh,
+                      drawHand + dh,
+                      -99999, 99999,
+                      (0,1,0))
+        else:
+            rectangle(drawHand - 1,drawHand + 1,
+                      43,45,(1,1,1))
+
+    return a
+
+

dreamcoder/dreamcoder.py

@@ -0,0 +1,1074 @@
+import datetime
+
+import dill
+
+from dreamcoder.compression import induceGrammar
+from dreamcoder.utilities import *
+try:
+    from dreamcoder.recognition import *
+except:
+    eprint("Failure loading recognition - only acceptable if using pypy ")
+from dreamcoder.enumeration import *
+from dreamcoder.fragmentGrammar import *
+from dreamcoder.taskBatcher import *
+from dreamcoder.primitiveGraph import graphPrimitives
+from dreamcoder.dreaming import backgroundHelmholtzEnumeration
+
+
+class ECResult():
+    def __init__(self, _=None,
+                 frontiersOverTime=None,
+                 testingSearchTime=None,
+                 learningCurve=None,
+                 grammars=None,
+                 taskSolutions=None,
+                 averageDescriptionLength=None,
+                 parameters=None,
+                 recognitionModel=None,
+                 searchTimes=None,
+                 recognitionTaskMetrics=None,
+                 numTestingTasks=None,
+                 sumMaxll=None,
+                 testingSumMaxll=None,
+                 hitsAtEachWake=None,
+                 timesAtEachWake=None,
+                 allFrontiers=None):
+        self.frontiersOverTime = {} # Map from task to [frontier at iteration 1, frontier at iteration 2, ...]
+        self.hitsAtEachWake = hitsAtEachWake or []
+        self.timesAtEachWake = timesAtEachWake or []
+        self.testingSearchTime = testingSearchTime or []
+        self.searchTimes = searchTimes or []
+        self.trainSearchTime = {} # map from task to search time
+        self.testSearchTime = {} # map from task to search time
+        self.recognitionTaskMetrics = recognitionTaskMetrics or {} 
+        self.recognitionModel = recognitionModel
+        self.averageDescriptionLength = averageDescriptionLength or []
+        self.parameters = parameters
+        self.learningCurve = learningCurve or []
+        self.grammars = grammars or []
+        self.taskSolutions = taskSolutions or {}
+        self.numTestingTasks = numTestingTasks
+        self.sumMaxll = sumMaxll or [] #TODO name change 
+        self.testingSumMaxll = testingSumMaxll or [] #TODO name change
+        self.allFrontiers = allFrontiers or {}
+
+    def __repr__(self):
+        attrs = ["{}={}".format(k, v) for k, v in self.__dict__.items()]
+        return "ECResult({})".format(", ".join(attrs))
+
+    def getTestingTasks(self):
+        testing = []
+        training = self.taskSolutions.keys()
+        for t in self.recognitionTaskMetrics:
+            if isinstance(t, Task) and t not in training: testing.append(t)
+        return testing
+
+
+    def recordFrontier(self, frontier):
+        t = frontier.task
+        if t not in self.frontiersOverTime: self.frontiersOverTime[t] = []
+        self.frontiersOverTime[t].append(frontier)
+
+    # Linux does not like files that have more than 256 characters
+    # So when exporting the results we abbreviate the parameters
+    abbreviations = {"frontierSize": "fs",
+                     "useDSL": "DSL",
+                     "taskReranker": "TRR",
+                     "matrixRank": "MR",
+                     "reuseRecognition": "RR",
+                     "ensembleSize": "ES",
+                     "recognitionTimeout": "RT",
+                     "recognitionSteps": "RS",
+                     "iterations": "it",
+                     "maximumFrontier": "MF",
+                     "pseudoCounts": "pc",
+                     "auxiliaryLoss": "aux",
+                     "structurePenalty": "L",
+                     "helmholtzRatio": "HR",
+                     "biasOptimal": "BO",
+                     "contextual": "CO",
+                     "topK": "K",
+                     "enumerationTimeout": "ET",
+                     "useRecognitionModel": "rec",
+                     "use_ll_cutoff": "llcut",
+                     "topk_use_only_likelihood": "topkNotMAP",
+                     "activation": "act",
+                     "storeTaskMetrics": 'STM',
+                     "topkNotMAP": "tknm",
+                     "rewriteTaskMetrics": "RW",
+                     'taskBatchSize': 'batch'}
+
+    @staticmethod
+    def abbreviate(parameter): return ECResult.abbreviations.get(parameter, parameter)
+
+    @staticmethod
+    def parameterOfAbbreviation(abbreviation):
+        return ECResult.abbreviationToParameter.get(abbreviation, abbreviation)
+
+    @staticmethod
+    def clearRecognitionModel(path):
+        SUFFIX = '.pickle'
+        assert path.endswith(SUFFIX)
+        
+        with open(path,'rb') as handle:
+            result = dill.load(handle)
+        
+        result.recognitionModel = None
+        
+        clearedPath = path[:-len(SUFFIX)] + "_graph=True" + SUFFIX
+        with open(clearedPath,'wb') as handle:
+            result = dill.dump(result, handle)
+        eprint(" [+] Cleared recognition model from:")
+        eprint("     %s"%path)
+        eprint("     and exported to:")
+        eprint("     %s"%clearedPath)
+        eprint("     Use this one for graphing.")
+
+
+ECResult.abbreviationToParameter = {
+    v: k for k, v in ECResult.abbreviations.items()}
+
+
+def explorationCompression(*arguments, **keywords):
+    for r in ecIterator(*arguments, **keywords):
+        pass
+    return r
+
+
+def ecIterator(grammar, tasks,
+               _=None,
+               useDSL=True,
+               noConsolidation=False,
+               mask=False,
+               seed=0,
+               addFullTaskMetrics=False,
+               matrixRank=None,
+               solver='ocaml',
+               compressor="rust",
+               biasOptimal=False,
+               contextual=False,
+               testingTasks=[],
+               iterations=None,
+               resume=None,
+               enumerationTimeout=None,
+               testingTimeout=None,
+               testEvery=1,
+               reuseRecognition=False,
+               ensembleSize=1,
+               useRecognitionModel=True,
+               recognitionTimeout=None,
+               recognitionSteps=None,
+               helmholtzRatio=0.,
+               featureExtractor=None,
+               activation='relu',
+               topK=1,
+               topk_use_only_likelihood=False,
+               use_map_search_times=True,
+               maximumFrontier=None,
+               pseudoCounts=1.0, aic=1.0,
+               structurePenalty=0.001, arity=0,
+               evaluationTimeout=1.0,  # seconds
+               taskBatchSize=None,
+               taskReranker='default',
+               CPUs=1,
+               cuda=False,
+               message="",
+               outputPrefix=None,
+               storeTaskMetrics=False,
+               rewriteTaskMetrics=True,
+               auxiliaryLoss=False,
+               custom_wake_generative=None):
+    if enumerationTimeout is None:
+        eprint(
+            "Please specify an enumeration timeout:",
+            "explorationCompression(..., enumerationTimeout = ..., ...)")
+        assert False
+    if iterations is None:
+        eprint(
+            "Please specify a iteration count: explorationCompression(..., iterations = ...)")
+        assert False
+    if useRecognitionModel and featureExtractor is None:
+        eprint("Warning: Recognition model needs feature extractor.",
+               "Ignoring recognition model.")
+        useRecognitionModel = False
+    if ensembleSize > 1 and not useRecognitionModel:
+        eprint("Warning: ensemble size requires using the recognition model, aborting.")
+        assert False
+    if biasOptimal and not useRecognitionModel:
+        eprint("Bias optimality only applies to recognition models, aborting.")
+        assert False
+    if contextual and not useRecognitionModel:
+        eprint("Contextual only applies to recognition models, aborting")
+        assert False
+    if reuseRecognition and not useRecognitionModel:
+        eprint("Reuse of recognition model weights at successive iteration only applies to recognition models, aborting")
+        assert False
+    if matrixRank is not None and not contextual:
+        eprint("Matrix rank only applies to contextual recognition models, aborting")
+        assert False
+    assert useDSL or useRecognitionModel, "You specified that you didn't want to use the DSL AND you don't want to use the recognition model. Figure out what you want to use."
+    if testingTimeout > 0 and len(testingTasks) == 0:
+        eprint("You specified a testingTimeout, but did not provide any held out testing tasks, aborting.")
+        assert False
+
+    # We save the parameters that were passed into EC
+    # This is for the purpose of exporting the results of the experiment
+    parameters = {
+        k: v for k,
+        v in locals().items() if k not in {
+            "tasks",
+            "use_map_search_times",
+            "seed",
+            "activation",
+            "grammar",
+            "cuda",
+            "_",
+            "testingTimeout",
+            "testEvery",
+            "message",
+            "CPUs",
+            "outputPrefix",
+            "resume",
+            "resumeFrontierSize",
+            "addFullTaskMetrics",
+            "featureExtractor",
+            "evaluationTimeout",
+            "testingTasks",
+            "compressor",
+            "custom_wake_generative"} and v is not None}
+    if not useRecognitionModel:
+        for k in {"helmholtzRatio", "recognitionTimeout", "biasOptimal", "mask",
+                  "contextual", "matrixRank", "reuseRecognition", "auxiliaryLoss", "ensembleSize"}:
+            if k in parameters: del parameters[k]
+    else: del parameters["useRecognitionModel"];
+    if useRecognitionModel and not contextual:
+        if "matrixRank" in parameters:
+            del parameters["matrixRank"]
+        if "mask" in parameters:
+            del parameters["mask"]
+    if not mask and 'mask' in parameters: del parameters["mask"]
+    if not auxiliaryLoss and 'auxiliaryLoss' in parameters: del parameters['auxiliaryLoss']
+    if not useDSL:
+        for k in {"structurePenalty", "pseudoCounts", "aic"}:
+            del parameters[k]
+    else: del parameters["useDSL"]
+    
+    # Uses `parameters` to construct the checkpoint path
+    def checkpointPath(iteration, extra=""):
+        parameters["iterations"] = iteration
+        kvs = [
+            "{}={}".format(
+                ECResult.abbreviate(k),
+                parameters[k]) for k in sorted(
+                parameters.keys())]
+        return "{}_{}{}.pickle".format(outputPrefix, "_".join(kvs), extra)
+
+    if message:
+        message = " (" + message + ")"
+    eprint("Running EC%s on %s @ %s with %d CPUs and parameters:" %
+           (message, os.uname()[1], datetime.datetime.now(), CPUs))
+    for k, v in parameters.items():
+        eprint("\t", k, " = ", v)
+    eprint("\t", "evaluationTimeout", " = ", evaluationTimeout)
+    eprint("\t", "cuda", " = ", cuda)
+    eprint()
+
+    if addFullTaskMetrics:
+        assert resume is not None, "--addFullTaskMetrics requires --resume"
+
+    def reportMemory():
+        eprint(f"Currently using this much memory: {getThisMemoryUsage()}")
+    
+    # Restore checkpoint
+    if resume is not None:
+        try:
+            resume = int(resume)
+            path = checkpointPath(resume)
+        except ValueError:
+            path = resume
+        with open(path, "rb") as handle:
+            result = dill.load(handle)
+        resume = len(result.grammars) - 1
+        eprint("Loaded checkpoint from", path)
+        grammar = result.grammars[-1] if result.grammars else grammar
+    else:  # Start from scratch
+        #for graphing of testing tasks
+        numTestingTasks = len(testingTasks) if len(testingTasks) != 0 else None
+
+        result = ECResult(parameters=parameters,            
+                          grammars=[grammar],
+                          taskSolutions={
+                              t: Frontier([],
+                                          task=t) for t in tasks},
+                          recognitionModel=None, numTestingTasks=numTestingTasks,
+                          allFrontiers={
+                              t: Frontier([],
+                                          task=t) for t in tasks})
+
+
+    # Set up the task batcher.
+    if taskReranker == 'default':
+        taskBatcher = DefaultTaskBatcher()
+    elif taskReranker == 'random':
+        taskBatcher = RandomTaskBatcher()
+    elif taskReranker == 'randomShuffle':
+        taskBatcher = RandomShuffleTaskBatcher(seed)
+    elif taskReranker == 'unsolved':
+        taskBatcher = UnsolvedTaskBatcher()
+    elif taskReranker == 'unsolvedEntropy':
+        taskBatcher = UnsolvedEntropyTaskBatcher()
+    elif taskReranker == 'unsolvedRandomEntropy':
+        taskBatcher = UnsolvedRandomEntropyTaskBatcher()
+    elif taskReranker == 'randomkNN':
+        taskBatcher = RandomkNNTaskBatcher()
+    elif taskReranker == 'randomLowEntropykNN':
+        taskBatcher = RandomLowEntropykNNTaskBatcher()
+    else:
+        eprint("Invalid task reranker: " + taskReranker + ", aborting.")
+        assert False
+
+    # Check if we are just updating the full task metrics
+    if addFullTaskMetrics:
+        if testingTimeout is not None and testingTimeout > enumerationTimeout:
+            enumerationTimeout = testingTimeout
+        if result.recognitionModel is not None:
+            _enumerator = lambda *args, **kw: result.recognitionModel.enumerateFrontiers(*args, **kw)
+        else: _enumerator = lambda *args, **kw: multicoreEnumeration(result.grammars[-1], *args, **kw)
+        enumerator = lambda *args, **kw: _enumerator(*args, 
+                                                     maximumFrontier=maximumFrontier, 
+                                                     CPUs=CPUs, evaluationTimeout=evaluationTimeout,
+                                                     solver=solver,
+                                                     **kw)
+        trainFrontiers, _, trainingTimes = enumerator(tasks, enumerationTimeout=enumerationTimeout)
+        testFrontiers, _, testingTimes = enumerator(testingTasks, enumerationTimeout=testingTimeout, testing=True)
+
+        recognizer = result.recognitionModel
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, trainingTimes, 'recognitionBestTimes')
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, recognizer.taskGrammarLogProductions(tasks), 'taskLogProductions')
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, recognizer.taskGrammarEntropies(tasks), 'taskGrammarEntropies')
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskAuxiliaryLossLayer(tasks), 'taskAuxiliaryLossLayer')
+        
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, testingTimes, 'heldoutTestingTimes')
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, recognizer.taskGrammarLogProductions(testingTasks), 'heldoutTaskLogProductions')
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, recognizer.taskGrammarEntropies(testingTasks), 'heldoutTaskGrammarEntropies')
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskAuxiliaryLossLayer(testingTasks), 'heldoutAuxiliaryLossLayer')
+
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, {f.task: f
+                                                                 for f in trainFrontiers + testFrontiers
+                                                                 if len(f) > 0},
+                                 'frontier')
+        SUFFIX = ".pickle"
+        assert path.endswith(SUFFIX)
+        path = path[:-len(SUFFIX)] + "_FTM=True" + SUFFIX
+        with open(path, "wb") as handle: dill.dump(result, handle)
+        if useRecognitionModel: ECResult.clearRecognitionModel(path)
+            
+        sys.exit(0)
+    
+    
+    for j in range(resume or 0, iterations):
+        if storeTaskMetrics and rewriteTaskMetrics:
+            eprint("Resetting task metrics for next iteration.")
+            result.recognitionTaskMetrics = {}
+
+        reportMemory()
+
+        # Evaluate on held out tasks if we have them
+        if testingTimeout > 0 and ((j % testEvery == 0) or (j == iterations - 1)):
+            eprint("Evaluating on held out testing tasks for iteration: %d" % (j))
+            evaluateOnTestingTasks(result, testingTasks, grammar,
+                                   CPUs=CPUs, maximumFrontier=maximumFrontier,
+                                   solver=solver,
+                                   enumerationTimeout=testingTimeout, evaluationTimeout=evaluationTimeout)            
+        # If we have to also enumerate Helmholtz frontiers,
+        # do this extra sneaky in the background
+        if useRecognitionModel and biasOptimal and helmholtzRatio > 0 and \
+           all( str(p) != "REAL" for p in grammar.primitives ): # real numbers don't support this
+            # the DSL is fixed, so the dreams are also fixed. don't recompute them.
+            if useDSL or 'helmholtzFrontiers' not in locals():
+                helmholtzFrontiers = backgroundHelmholtzEnumeration(tasks, grammar, enumerationTimeout,
+                                                                    evaluationTimeout=evaluationTimeout,
+                                                                    special=featureExtractor.special)
+            else:
+                print("Reusing dreams from previous iteration.")
+        else:
+            helmholtzFrontiers = lambda: []
+
+        reportMemory()
+
+        # Get waking task batch.
+        wakingTaskBatch = taskBatcher.getTaskBatch(result, tasks, taskBatchSize, j)
+        eprint("Using a waking task batch of size: " + str(len(wakingTaskBatch)))
+
+        # WAKING UP
+        if useDSL:
+            wake_generative = custom_wake_generative if custom_wake_generative is not None else default_wake_generative
+            topDownFrontiers, times = wake_generative(grammar, wakingTaskBatch,
+                                                      solver=solver,
+                                                      maximumFrontier=maximumFrontier,
+                                                      enumerationTimeout=enumerationTimeout,
+                                                      CPUs=CPUs,
+                                                      evaluationTimeout=evaluationTimeout)
+            result.trainSearchTime = {t: tm for t, tm in times.items() if tm is not None}
+        else:
+            eprint("Skipping top-down enumeration because we are not using the generative model")
+            topDownFrontiers, times = [], {t: None for t in wakingTaskBatch }
+
+        tasksHitTopDown = {f.task for f in topDownFrontiers if not f.empty}
+        result.hitsAtEachWake.append(len(tasksHitTopDown))
+
+        reportMemory()
+
+        # Combine topDownFrontiers from this task batch with all frontiers.
+        for f in topDownFrontiers:
+            if f.task not in result.allFrontiers: continue # backward compatibility with old checkpoints
+            result.allFrontiers[f.task] = result.allFrontiers[f.task].combine(f).topK(maximumFrontier)
+
+        eprint("Frontiers discovered top down: " + str(len(tasksHitTopDown)))
+        eprint("Total frontiers: " + str(len([f for f in result.allFrontiers.values() if not f.empty])))
+
+        # Train + use recognition model
+        if useRecognitionModel:
+            # Should we initialize the weights to be what they were before?
+            previousRecognitionModel = None
+            if reuseRecognition and result.recognitionModel is not None:
+                previousRecognitionModel = result.recognitionModel
+
+            thisRatio = helmholtzRatio
+            #if j == 0 and not biasOptimal: thisRatio = 0
+            if all( f.empty for f in result.allFrontiers.values() ): thisRatio = 1.                
+
+            tasksHitBottomUp = \
+             sleep_recognition(result, grammar, wakingTaskBatch, tasks, testingTasks, result.allFrontiers.values(),
+                               ensembleSize=ensembleSize, featureExtractor=featureExtractor, mask=mask,
+                               activation=activation, contextual=contextual, biasOptimal=biasOptimal,
+                               previousRecognitionModel=previousRecognitionModel, matrixRank=matrixRank,
+                               timeout=recognitionTimeout, evaluationTimeout=evaluationTimeout,
+                               enumerationTimeout=enumerationTimeout,
+                               helmholtzRatio=thisRatio, helmholtzFrontiers=helmholtzFrontiers(),
+                               auxiliaryLoss=auxiliaryLoss, cuda=cuda, CPUs=CPUs, solver=solver,
+                               recognitionSteps=recognitionSteps, maximumFrontier=maximumFrontier)
+
+            showHitMatrix(tasksHitTopDown, tasksHitBottomUp, wakingTaskBatch)
+            
+        # Record the new topK solutions
+        result.taskSolutions = {f.task: f.topK(topK)
+                                for f in result.allFrontiers.values()}
+        for f in result.allFrontiers.values(): result.recordFrontier(f)
+        result.learningCurve += [
+            sum(f is not None and not f.empty for f in result.taskSolutions.values())]
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, {f.task: f
+                                                                 for f in result.allFrontiers.values()
+                                                                 if len(f) > 0},
+                                 'frontier')                
+        
+        # Sleep-G
+        if useDSL and not(noConsolidation):
+            eprint(f"Currently using this much memory: {getThisMemoryUsage()}")
+            grammar = consolidate(result, grammar, topK=topK, pseudoCounts=pseudoCounts, arity=arity, aic=aic,
+                                  structurePenalty=structurePenalty, compressor=compressor, CPUs=CPUs,
+                                  iteration=j)
+            eprint(f"Currently using this much memory: {getThisMemoryUsage()}")
+        else:
+            eprint("Skipping consolidation.")
+            result.grammars.append(grammar)
+            
+        if outputPrefix is not None:
+            path = checkpointPath(j + 1)
+            with open(path, "wb") as handle:
+                try:
+                    dill.dump(result, handle)
+                except TypeError as e:
+                    eprint(result)
+                    assert(False)
+            eprint("Exported checkpoint to", path)
+            if useRecognitionModel:
+                ECResult.clearRecognitionModel(path)
+
+            graphPrimitives(result, "%s_primitives_%d_"%(outputPrefix,j))
+            
+
+        yield result
+
+
+def showHitMatrix(top, bottom, tasks):
+    tasks = set(tasks)
+
+    total = bottom | top
+    eprint(len(total), "/", len(tasks), "total hit tasks")
+    bottomMiss = tasks - bottom
+    topMiss = tasks - top
+
+    eprint("{: <13s}{: ^13s}{: ^13s}".format("", "bottom miss", "bottom hit"))
+    eprint("{: <13s}{: ^13d}{: ^13d}".format("top miss",
+                                             len(bottomMiss & topMiss),
+                                             len(bottom & topMiss)))
+    eprint("{: <13s}{: ^13d}{: ^13d}".format("top hit",
+                                             len(top & bottomMiss),
+                                             len(top & bottom)))
+
+def evaluateOnTestingTasks(result, testingTasks, grammar, _=None,
+                           CPUs=None, solver=None, maximumFrontier=None, enumerationTimeout=None, evaluationTimeout=None):
+    if result.recognitionModel is not None:
+        recognizer = result.recognitionModel
+        testingFrontiers, times = \
+         recognizer.enumerateFrontiers(testingTasks, 
+                                       CPUs=CPUs,
+                                       solver=solver,
+                                       maximumFrontier=maximumFrontier,
+                                       enumerationTimeout=enumerationTimeout,
+                                       evaluationTimeout=evaluationTimeout,
+                                       testing=True)
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, recognizer.taskGrammarLogProductions(testingTasks), 'heldoutTaskLogProductions')
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, recognizer.taskGrammarEntropies(testingTasks), 'heldoutTaskGrammarEntropies')
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, recognizer.taskGrammarEntropies(testingTasks), 'heldoutTaskGrammarEntropies')
+    else:
+        testingFrontiers, times = multicoreEnumeration(grammar, testingTasks, 
+                                                       solver=solver,
+                                                       maximumFrontier=maximumFrontier,
+                                                       enumerationTimeout=enumerationTimeout,
+                                                       CPUs=CPUs,
+                                                       evaluationTimeout=evaluationTimeout,
+                                                       testing=True)
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, times, 'heldoutTestingTimes')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics,
+                                     {f.task: f for f in testingFrontiers if len(f) > 0 },
+                                     'frontier')
+    for f in testingFrontiers: result.recordFrontier(f)
+    result.testSearchTime = {t: tm for t, tm in times.items() if tm is not None}
+    times = [t for t in times.values() if t is not None ]
+    eprint("\n".join(f.summarize() for f in testingFrontiers))
+    summaryStatistics("Testing tasks", times)
+    eprint("Hits %d/%d testing tasks" % (len(times), len(testingTasks)))
+    result.testingSearchTime.append(times)
+
+        
+def default_wake_generative(grammar, tasks, 
+                    maximumFrontier=None,
+                    enumerationTimeout=None,
+                    CPUs=None,
+                    solver=None,
+                    evaluationTimeout=None):
+    topDownFrontiers, times = multicoreEnumeration(grammar, tasks, 
+                                                   maximumFrontier=maximumFrontier,
+                                                   enumerationTimeout=enumerationTimeout,
+                                                   CPUs=CPUs,
+                                                   solver=solver,
+                                                   evaluationTimeout=evaluationTimeout)
+    eprint("Generative model enumeration results:")
+    eprint(Frontier.describe(topDownFrontiers))
+    summaryStatistics("Generative model", [t for t in times.values() if t is not None])
+    return topDownFrontiers, times
+
+def sleep_recognition(result, grammar, taskBatch, tasks, testingTasks, allFrontiers, _=None,
+                      ensembleSize=1, featureExtractor=None, matrixRank=None, mask=False,
+                      activation=None, contextual=True, biasOptimal=True,
+                      previousRecognitionModel=None, recognitionSteps=None,
+                      timeout=None, enumerationTimeout=None, evaluationTimeout=None,
+                      helmholtzRatio=None, helmholtzFrontiers=None, maximumFrontier=None,
+                      auxiliaryLoss=None, cuda=None, CPUs=None, solver=None):
+    eprint("Using an ensemble size of %d. Note that we will only store and test on the best recognition model." % ensembleSize)
+
+    featureExtractorObjects = [featureExtractor(tasks, testingTasks=testingTasks, cuda=cuda) for i in range(ensembleSize)]
+    recognizers = [RecognitionModel(featureExtractorObjects[i],
+                                    grammar,
+                                    mask=mask,
+                                    rank=matrixRank,
+                                    activation=activation,
+                                    cuda=cuda,
+                                    contextual=contextual,
+                                    previousRecognitionModel=previousRecognitionModel,
+                                    id=i) for i in range(ensembleSize)]
+    eprint(f"Currently using this much memory: {getThisMemoryUsage()}")
+    trainedRecognizers = parallelMap(min(CPUs,len(recognizers)),
+                                     lambda recognizer: recognizer.train(allFrontiers,
+                                                                         biasOptimal=biasOptimal,
+                                                                         helmholtzFrontiers=helmholtzFrontiers, 
+                                                                         CPUs=CPUs,
+                                                                         evaluationTimeout=evaluationTimeout,
+                                                                         timeout=timeout,
+                                                                         steps=recognitionSteps,
+                                                                         helmholtzRatio=helmholtzRatio,
+                                                                         auxLoss=auxiliaryLoss,
+                                                                         vectorized=True),
+                                     recognizers,
+                                     seedRandom=True)
+    eprint(f"Currently using this much memory: {getThisMemoryUsage()}")
+    # Enumerate frontiers for each of the recognizers.
+    eprint("Trained an ensemble of %d recognition models, now enumerating." % len(trainedRecognizers))
+    ensembleFrontiers, ensembleTimes, ensembleRecognitionTimes = [], [], []
+    mostTasks = 0
+    bestRecognizer = None
+    totalTasksHitBottomUp = set()
+    for recIndex, recognizer in enumerate(trainedRecognizers):
+        eprint("Enumerating from recognizer %d of %d" % (recIndex, len(trainedRecognizers)))
+        bottomupFrontiers, allRecognitionTimes = \
+                        recognizer.enumerateFrontiers(taskBatch, 
+                                                      CPUs=CPUs,
+                                                      maximumFrontier=maximumFrontier,
+                                                      enumerationTimeout=enumerationTimeout,
+                                                      evaluationTimeout=evaluationTimeout,
+                                                      solver=solver)
+        ensembleFrontiers.append(bottomupFrontiers)
+        ensembleTimes.append([t for t in allRecognitionTimes.values() if t is not None])
+        ensembleRecognitionTimes.append(allRecognitionTimes)
+
+        recognizerTasksHitBottomUp = {f.task for f in bottomupFrontiers if not f.empty}
+        totalTasksHitBottomUp.update(recognizerTasksHitBottomUp)
+        eprint("Recognizer %d solved %d/%d tasks; total tasks solved is now %d." % (recIndex, len(recognizerTasksHitBottomUp), len(tasks), len(totalTasksHitBottomUp)))
+        if len(recognizerTasksHitBottomUp) >= mostTasks:
+            # TODO (cathywong): could consider keeping the one that put the highest likelihood on the solved tasks.
+            bestRecognizer = recIndex
+
+    # Store the recognizer that discovers the most frontiers in the result.
+    eprint("Best recognizer: %d." % bestRecognizer)
+    result.recognitionModel = trainedRecognizers[bestRecognizer]
+    result.trainSearchTime = {tk: tm for tk, tm in ensembleRecognitionTimes[bestRecognizer].items()
+                              if tm is not None}
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, ensembleRecognitionTimes[bestRecognizer], 'recognitionBestTimes')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskHiddenStates(tasks), 'hiddenState')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskGrammarLogProductions(tasks), 'taskLogProductions')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskGrammarEntropies(tasks), 'taskGrammarEntropies')
+    if contextual:
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics,
+                                 result.recognitionModel.taskGrammarStartProductions(tasks),
+                                 'startProductions')
+
+    result.hitsAtEachWake.append(len(totalTasksHitBottomUp))
+    eprint(f"Currently using this much memory: {getThisMemoryUsage()}")
+
+    """ Rescore and combine the frontiers across the ensemble of recognition models."""
+    eprint("Recognition model enumeration results for the best recognizer.")
+    eprint(Frontier.describe(ensembleFrontiers[bestRecognizer]))
+    summaryStatistics("Recognition model", ensembleTimes[bestRecognizer])
+
+    eprint("Cumulative results for the full ensemble of %d recognizers: " % len(trainedRecognizers))
+    # Rescore all of the ensemble frontiers according to the generative model
+    # and then combine w/ original frontiers
+    for bottomupFrontiers in ensembleFrontiers:
+        for b in bottomupFrontiers:
+            if b.task not in result.allFrontiers: continue # backwards compatibility with old checkpoints
+            result.allFrontiers[b.task] = result.allFrontiers[b.task].\
+                                          combine(grammar.rescoreFrontier(b)).\
+                                          topK(maximumFrontier)
+
+    eprint("Frontiers discovered bottom up: " + str(len(totalTasksHitBottomUp)))
+    eprint("Total frontiers: " + str(len([f for f in result.allFrontiers.values() if not f.empty])))
+
+    result.searchTimes.append(ensembleTimes[bestRecognizer])
+    if len(ensembleTimes[bestRecognizer]) > 0:
+        eprint("Average search time: ", int(mean(ensembleTimes[bestRecognizer]) + 0.5),
+               "sec.\tmedian:", int(median(ensembleTimes[bestRecognizer]) + 0.5),
+               "\tmax:", int(max(ensembleTimes[bestRecognizer]) + 0.5),
+               "\tstandard deviation", int(standardDeviation(ensembleTimes[bestRecognizer]) + 0.5))
+    return totalTasksHitBottomUp
+
+def consolidate(result, grammar, _=None, topK=None, arity=None, pseudoCounts=None, aic=None,
+                structurePenalty=None, compressor=None, CPUs=None, iteration=None):
+    eprint("Showing the top 5 programs in each frontier being sent to the compressor:")
+    for f in result.allFrontiers.values():
+        if f.empty:
+            continue
+        eprint(f.task)
+        for e in f.normalize().topK(5):
+            eprint("%.02f\t%s" % (e.logPosterior, e.program))
+        eprint()
+
+    # First check if we have supervision at the program level for any task that was not solved
+    needToSupervise = {f.task for f in result.allFrontiers.values()
+                       if f.task.supervision is not None and f.empty}
+    compressionFrontiers = [f.replaceWithSupervised(grammar) if f.task in needToSupervise else f
+                            for f in result.allFrontiers.values() ]
+
+    if len([f for f in compressionFrontiers if not f.empty]) == 0:
+        eprint("No compression frontiers; not inducing a grammar this iteration.")
+    else:
+        grammar, compressionFrontiers = induceGrammar(grammar, compressionFrontiers,
+                                                      topK=topK,
+                                                      pseudoCounts=pseudoCounts, a=arity,
+                                                      aic=aic, structurePenalty=structurePenalty,
+                                                      topk_use_only_likelihood=False,
+                                                      backend=compressor, CPUs=CPUs, iteration=iteration)
+        # Store compression frontiers in the result.
+        for c in compressionFrontiers:
+            result.allFrontiers[c.task] = c.topK(0) if c in needToSupervise else c
+
+
+    result.grammars.append(grammar)
+    eprint("Grammar after iteration %d:" % (iteration + 1))
+    eprint(grammar)
+    
+    return grammar
+    
+
+
+def commandlineArguments(_=None,
+                         iterations=None,
+                         enumerationTimeout=None,
+                         testEvery=1,
+                         topK=1,
+                         reuseRecognition=False,
+                         CPUs=1,
+                         solver='ocaml',
+                         compressor="ocaml",
+                         useRecognitionModel=True,
+                         recognitionTimeout=None,
+                         activation='relu',
+                         helmholtzRatio=1.,
+                         featureExtractor=None,
+                         cuda=None,
+                         maximumFrontier=None,
+                         pseudoCounts=1.0, aic=1.0,
+                         structurePenalty=0.001, a=0,
+                         taskBatchSize=None, taskReranker="default",
+                         extras=None,
+                         storeTaskMetrics=False,
+                        rewriteTaskMetrics=True):
+    if cuda is None:
+        cuda = torch.cuda.is_available()
+    print("CUDA is available?:", torch.cuda.is_available())
+    print("using cuda?:", cuda)
+    import argparse
+    parser = argparse.ArgumentParser(description="")
+    parser.add_argument("--resume",
+                        help="Resumes EC algorithm from checkpoint. You can either pass in the path of a checkpoint, or you can pass in the iteration to resume from, in which case it will try to figure out the path.",
+                        default=None,
+                        type=str)
+    parser.add_argument("-i", "--iterations",
+                        help="default: %d" % iterations,
+                        default=iterations,
+                        type=int)
+    parser.add_argument("-t", "--enumerationTimeout",
+                        default=enumerationTimeout,
+                        help="In seconds. default: %s" % enumerationTimeout,
+                        type=int)
+    parser.add_argument("-R", "--recognitionTimeout",
+                        default=recognitionTimeout,
+                        help="In seconds. Amount of time to train the recognition model on each iteration. Defaults to enumeration timeout.",
+                        type=int)
+    parser.add_argument("-RS", "--recognitionSteps",
+                        default=None,
+                        help="Number of gradient steps to train the recognition model. Can be specified instead of train time.",
+                        type=int)
+    parser.add_argument(
+        "-k",
+        "--topK",
+        default=topK,
+        help="When training generative and discriminative models, we train them to fit the top K programs. Ideally we would train them to fit the entire frontier, but this is often intractable. default: %d" %
+        topK,
+        type=int)
+    parser.add_argument("-p", "--pseudoCounts",
+                        default=pseudoCounts,
+                        help="default: %f" % pseudoCounts,
+                        type=float)
+    parser.add_argument("-b", "--aic",
+                        default=aic,
+                        help="default: %f" % aic,
+                        type=float)
+    parser.add_argument("-l", "--structurePenalty",
+                        default=structurePenalty,
+                        help="default: %f" % structurePenalty,
+                        type=float)
+    parser.add_argument("-a", "--arity",
+                        default=a,
+                        help="default: %d" % a,
+                        type=int)
+    parser.add_argument("-c", "--CPUs",
+                        default=CPUs,
+                        help="default: %d" % CPUs,
+                        type=int)
+    parser.add_argument("--no-cuda",
+                        action="store_false",
+                        dest="cuda",
+                        help="""cuda will be used if available (which it %s),
+                        unless this is set""" % ("IS" if cuda else "ISN'T"))
+    parser.add_argument("-m", "--maximumFrontier",
+                        help="""Even though we enumerate --frontierSize
+                        programs, we might want to only keep around the very
+                        best for performance reasons. This is a cut off on the
+                        maximum size of the frontier that is kept around.
+                        Default: %s""" % maximumFrontier,
+                        type=int)
+    parser.add_argument("--reuseRecognition",
+                        help="""Should we initialize recognition model weights to be what they were at the previous DreamCoder iteration? Default: %s""" % reuseRecognition,
+                        default=reuseRecognition,
+                        action="store_true")
+    parser.add_argument("--recognition",
+                        dest="useRecognitionModel",
+                        action="store_true",
+                        help="""Enable bottom-up neural recognition model.
+                        Default: %s""" % useRecognitionModel)
+    parser.add_argument("--ensembleSize",
+                        dest="ensembleSize",
+                        default=1,
+                        help="Number of recognition models to train and enumerate from at each iteration.",
+                        type=int)
+    parser.add_argument("-g", "--no-recognition",
+                        dest="useRecognitionModel",
+                        action="store_false",
+                        help="""Disable bottom-up neural recognition model.
+                        Default: %s""" % (not useRecognitionModel))
+    parser.add_argument("-d", "--no-dsl",
+                        dest="useDSL",
+                        action="store_false",
+                        help="""Disable DSL enumeration and updating.""")
+    parser.add_argument("--no-consolidation",
+                        dest="noConsolidation",
+                        action="store_true",
+                        help="""Disable DSL updating.""")
+    parser.add_argument(
+        "--testingTimeout",
+        type=int,
+        dest="testingTimeout",
+        default=0,
+        help="Number of seconds we should spend evaluating on each held out testing task.")
+    parser.add_argument(
+        "--testEvery",
+        type=int,
+        dest="testEvery",
+        default=1,
+        help="Run heldout testing every X iterations."
+        )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="Random seed. Currently this only matters for random batching strategies.")
+    parser.add_argument(
+        "--activation",
+        choices=[
+            "relu",
+            "sigmoid",
+            "tanh"],
+        default=activation,
+        help="""Activation function for neural recognition model.
+                        Default: %s""" %
+        activation)
+    parser.add_argument(
+        "--solver",
+        choices=[
+            "ocaml",
+            "pypy",
+            "python"],
+        default=solver,
+        help="""Solver for enumeration.
+                        Default: %s""" %
+        solver)
+    parser.add_argument(
+        "-r",
+        "--Helmholtz",
+        dest="helmholtzRatio",
+        help="""When training recognition models, what fraction of the training data should be samples from the generative model? Default %f""" %
+        helmholtzRatio,
+        default=helmholtzRatio,
+        type=float)
+    parser.add_argument(
+        "--compressor",
+        default=compressor,
+        choices=["pypy","rust","vs","pypy_vs","ocaml","memorize"])
+    parser.add_argument(
+        "--matrixRank",
+        help="Maximum rank of bigram transition matrix for contextual recognition model. Defaults to full rank.",
+        default=None,
+        type=int)
+    parser.add_argument(
+        "--mask",
+        help="Unconditional bigram masking",
+        default=False, action="store_true")
+    parser.add_argument("--biasOptimal",
+                        help="Enumerate dreams rather than sample them & bias-optimal recognition objective",
+                        default=False, action="store_true")
+    parser.add_argument("--contextual",
+                        help="bigram recognition model (default is unigram model)",
+                        default=False, action="store_true")
+    parser.add_argument("--clear-recognition",
+                        dest="clear-recognition",
+                        help="Clears the recognition model from a checkpoint. Necessary for graphing results with recognition models, because pickle is kind of stupid sometimes.",
+                        default=None,
+                        type=str)
+    parser.add_argument("--primitive-graph",
+                        dest="primitive-graph",
+                        nargs='+',
+                        help="Displays a dependency graph of the learned primitives",
+                        default=None,
+                        type=str)
+    parser.add_argument(
+        "--taskBatchSize",
+        dest="taskBatchSize",
+        help="Number of tasks to train on during wake. Defaults to all tasks if None.",
+        default=None,
+        type=int)
+    parser.add_argument(
+        "--taskReranker",
+        dest="taskReranker",
+        help="Reranking function used to order the tasks we train on during waking.",
+        choices=[
+            "default",
+            "random",
+            "randomShuffle",
+            "unsolved",
+            "unsolvedEntropy",
+            "unsolvedRandomEntropy",
+            "randomkNN",
+            "randomLowEntropykNN"],
+        default=taskReranker,
+        type=str)
+    parser.add_argument(
+        "--storeTaskMetrics",
+        dest="storeTaskMetrics",
+        default=True,
+        help="Whether to store task metrics directly in the ECResults.",
+        action="store_true"
+        )
+    parser.add_argument(
+        "--rewriteTaskMetrics",
+        dest="rewriteTaskMetrics",
+        help="Whether to rewrite a new task metrics dictionary at each iteration, rather than retaining the old.",
+        action="store_true"
+        )
+    parser.add_argument("--addTaskMetrics",
+                        dest="addTaskMetrics",
+                        help="Creates a checkpoint with task metrics and no recognition model for graphing.",
+                        default=None,
+                        nargs='+',
+                        type=str)
+    parser.add_argument("--auxiliary",
+                        action="store_true", default=False,
+                        help="Add auxiliary classification loss to recognition network training",
+                        dest="auxiliaryLoss")
+    parser.add_argument("--addFullTaskMetrics",
+                        help="Only to be used in conjunction with --resume. Loads checkpoint, solves both testing and training tasks, stores frontiers, solve times, and task metrics, and then dies.",
+                        default=False,
+                        action="store_true")
+    parser.add_argument("--countParameters",
+                        help="Load a checkpoint then report how many parameters are in the recognition model.",
+                        default=None, type=str)
+    parser.set_defaults(useRecognitionModel=useRecognitionModel,
+                        useDSL=True,
+                        featureExtractor=featureExtractor,
+                        maximumFrontier=maximumFrontier,
+                        cuda=cuda)
+    if extras is not None:
+        extras(parser)
+    v = vars(parser.parse_args())
+    if v["clear-recognition"] is not None:
+        ECResult.clearRecognitionModel(v["clear-recognition"])
+        sys.exit(0)
+    else:
+        del v["clear-recognition"]
+        
+    if v["primitive-graph"] is not None:
+        
+        for n,pg in enumerate(v["primitive-graph"]):
+            with open(pg,'rb') as handle:
+                result = dill.load(handle)
+            graphPrimitives(result,f"figures/deepProgramLearning/{sys.argv[0]}{n}",view=True)
+        sys.exit(0)
+    else:
+        del v["primitive-graph"]
+
+    if v["addTaskMetrics"] is not None:
+        for path in v["addTaskMetrics"]:
+            with open(path,'rb') as handle:
+                result = dill.load(handle)
+            addTaskMetrics(result, path)
+        sys.exit(0)
+    else:
+        del v["addTaskMetrics"]
+
+    if v["useRecognitionModel"] and v["recognitionTimeout"] is None:
+        v["recognitionTimeout"] = v["enumerationTimeout"]
+
+    if v["countParameters"]:
+        with open(v["countParameters"], "rb") as handle:
+            result = dill.load(handle)
+        eprint("The recognition model has",
+               sum(p.numel() for p in result.recognitionModel.parameters() if p.requires_grad),
+               "trainable parameters and",
+               sum(p.numel() for p in result.recognitionModel.parameters() ),
+               "total parameters.\n",
+               "The feature extractor accounts for",
+               sum(p.numel() for p in result.recognitionModel.featureExtractor.parameters() ),
+               "of those parameters.\n",
+               "The grammar builder accounts for",
+               sum(p.numel() for p in result.recognitionModel.grammarBuilder.parameters() ),
+               "of those parameters.\n")
+        sys.exit(0)
+    del v["countParameters"]
+        
+        
+    return v
+
+def addTaskMetrics(result, path):
+    """Adds a task metrics to ECResults that were pickled without them."""
+    with torch.no_grad(): return addTaskMetrics_(result, path)
+def addTaskMetrics_(result, path):
+    SUFFIX = '.pickle'
+    assert path.endswith(SUFFIX)
+
+    tasks = result.taskSolutions.keys()
+    everyTask = set(tasks)
+    for t in result.recognitionTaskMetrics:
+        if isinstance(t, Task) and t not in everyTask: everyTask.add(t)
+
+    eprint(f"Found {len(tasks)} training tasks.")
+    eprint(f"Scrounged up {len(everyTask) - len(tasks)} testing tasks.")
+    if not hasattr(result, "recognitionTaskMetrics") or result.recognitionTaskMetrics is None:
+        result.recognitionTaskMetrics = {}
+
+    # If task has images, store them.
+    if hasattr(list(tasks)[0], 'getImage'):
+        images = {t: t.getImage(pretty=True) for t in tasks}
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, images, 'taskImages')
+
+    if hasattr(list(tasks)[0], 'highresolution'):
+        images = {t: t.highresolution for t in tasks}
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics, images, 'taskImages')
+
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.auxiliaryPrimitiveEmbeddings(), 'auxiliaryPrimitiveEmbeddings')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskAuxiliaryLossLayer(tasks), 'taskAuxiliaryLossLayer')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskAuxiliaryLossLayer(everyTask), 'every_auxiliaryLossLayer')
+
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskGrammarFeatureLogProductions(tasks), 'grammarFeatureLogProductions')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskGrammarFeatureLogProductions(everyTask), 'every_grammarFeatureLogProductions')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskGrammarLogProductions(tasks), 'contextualLogProductions')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskGrammarLogProductions(everyTask), 'every_contextualLogProductions')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskHiddenStates(tasks), 'hiddenState')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskHiddenStates(everyTask), 'every_hiddenState')
+    g = result.grammars[-2] # the final entry in result.grammars is a grammar that we have not used yet
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, {f.task: f.expectedProductionUses(g)
+                                                             for f in result.taskSolutions.values()
+                                                             if len(f) > 0},
+                             'expectedProductionUses')
+    updateTaskSummaryMetrics(result.recognitionTaskMetrics, {f.task: f.expectedProductionUses(g)
+                                                             for t, metrics in result.recognitionTaskMetrics.items()
+                                                             if "frontier" in metrics
+                                                             for f in [metrics["frontier"]] 
+                                                             if len(f) > 0},
+                             'every_expectedProductionUses')
+    if False:
+        eprint(f"About to do an expensive Monte Carlo simulation w/ {len(tasks)} tasks")
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics,
+                                 {task: result.recognitionModel.grammarOfTask(task).untorch().expectedUsesMonteCarlo(task.request, debug=False)
+                                  for task in tasks },
+                                 'expectedProductionUsesMonteCarlo')
+    try:
+        updateTaskSummaryMetrics(result.recognitionTaskMetrics,
+                                 result.recognitionModel.taskGrammarStartProductions(tasks),
+                                 'startProductions')
+    except: pass # can fail if we do not have a contextual model
+
+    #updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskGrammarLogProductions(tasks), 'task_no_parent_log_productions')
+    #updateTaskSummaryMetrics(result.recognitionTaskMetrics, result.recognitionModel.taskGrammarEntropies(tasks), 'taskGrammarEntropies')
+
+    result.recognitionModel = None
+        
+    clearedPath = path[:-len(SUFFIX)] + "_graph=True" + SUFFIX
+    with open(clearedPath,'wb') as handle:
+        result = dill.dump(result, handle)
+    eprint(" [+] Cleared recognition model from:")
+    eprint("     %s"%path)
+    eprint("     and exported to:")
+    eprint("     %s"%clearedPath)
+    eprint("     Use this one for graphing.")
+

dreamcoder/dreaming.py

@@ -0,0 +1,90 @@
+import json
+import os
+import subprocess
+
+from dreamcoder.domains.arithmetic.arithmeticPrimitives import k1, k0, addition, subtraction, multiplication
+from dreamcoder.frontier import Frontier, FrontierEntry
+from dreamcoder.grammar import Grammar
+from dreamcoder.program import Program
+from dreamcoder.task import Task
+from dreamcoder.type import arrow, tint
+from dreamcoder.utilities import tuplify, timing, eprint, get_root_dir, mean
+
+
+def helmholtzEnumeration(g, request, inputs, timeout, _=None,
+                         special=None, evaluationTimeout=None):
+    """Returns json (as text)"""
+    message = {"request": request.json(),
+               "timeout": timeout,
+               "DSL": g.json(),
+               "extras": inputs}
+    if evaluationTimeout: message["evaluationTimeout"] = evaluationTimeout
+    if special: message["special"] = special
+    message = json.dumps(message)
+    with open('/tmp/hm', 'w') as handle:
+        handle.write(message)
+    try:
+        binary = os.path.join(get_root_dir(), 'helmholtz')
+        process = subprocess.Popen(binary,
+                                   stdin=subprocess.PIPE,
+                                   stdout=subprocess.PIPE)
+        response, error = process.communicate(bytes(message, encoding="utf-8"))
+    except OSError as exc:
+        raise exc
+    return response
+
+
+def backgroundHelmholtzEnumeration(tasks, g, timeout, _=None,
+                                   special=None, evaluationTimeout=None):
+    from pathos.multiprocessing import Pool
+    requests = list({t.request for t in tasks})
+    inputs = {r: list({tuplify(xs)
+                       for t in tasks if t.request == r
+                       for xs, y in t.examples})
+              for r in requests}
+    workers = Pool(len(requests))
+    promises = [workers.apply_async(helmholtzEnumeration,
+                                    args=(g, r, inputs[r], float(timeout)),
+                                    kwds={'special': special,
+                                          'evaluationTimeout': evaluationTimeout})
+                for r in requests]
+
+    def get():
+        results = [p.get() for p in promises]
+        frontiers = []
+        with timing("(Helmholtz enumeration) Decoded json into frontiers"):
+            for request, result in zip(requests, results):
+                response = json.loads(result.decode("utf-8"))
+                for b, entry in enumerate(response):
+                    frontiers.append(Frontier([FrontierEntry(program=Program.parse(p),
+                                                             logPrior=entry["ll"],
+                                                             logLikelihood=0.)
+                                               for p in entry["programs"]],
+                                              task=Task(str(b),
+                                                        request,
+                                                        [])))
+        eprint("Total number of Helmholtz frontiers:", len(frontiers))
+        return frontiers
+
+    return get
+
+
+if __name__ == "__main__":
+    from dreamcoder.recognition import RecognitionModel, DummyFeatureExtractor
+    g = Grammar.uniform([k1, k0, addition, subtraction, multiplication])
+    frontiers = helmholtzEnumeration(g,
+                                     arrow(tint, tint),
+                                     [[0], [1], [2]],
+                                     10.)
+    eprint("average frontier size", mean(len(f.entries) for f in frontiers))
+    f = DummyFeatureExtractor([])
+    r = RecognitionModel(f, g, hidden=[], contextual=True)
+    r.trainBiasOptimal(frontiers, frontiers, steps=70)
+    g = r.grammarOfTask(frontiers[0].task).untorch()
+    frontiers = helmholtzEnumeration(g,
+                                     arrow(tint, tint),
+                                     [[0], [1], [2]],
+                                     10.)
+    for f in frontiers:
+        eprint(f.summarizeFull())
+    eprint("average frontier size", mean(len(f.entries) for f in frontiers))

dreamcoder/ec.py

@@ -0,0 +1,3 @@
+print("DEPRECATION NOTICE: this module (dreamcoder.ec) will be deleted soon, "
+      "please update your code to import from dreamcoder.dreamcoder instead")
+from dreamcoder.dreamcoder import *  # noqa

dreamcoder/enumeration.py

@@ -0,0 +1,469 @@
+from dreamcoder.likelihoodModel import AllOrNothingLikelihoodModel
+from dreamcoder.grammar import *
+from dreamcoder.utilities import get_root_dir
+
+import os
+import traceback
+import subprocess
+
+
+def multicoreEnumeration(g, tasks, _=None,
+                         enumerationTimeout=None,
+                         solver='ocaml',
+                         CPUs=1,
+                         maximumFrontier=None,
+                         verbose=True,
+                         evaluationTimeout=None,
+                         testing=False):
+    '''g: Either a Grammar, or a map from task to grammar.
+    Returns (list-of-frontiers, map-from-task-to-search-time)'''
+
+    # We don't use actual threads but instead use the multiprocessing
+    # library. This is because we need to be able to kill workers.
+    #from multiprocess import Process, Queue
+
+    from multiprocessing import Queue
+
+     # everything that gets sent between processes will be dilled
+    import dill
+
+    solvers = {"ocaml": solveForTask_ocaml,   
+               "pypy": solveForTask_pypy,   
+               "python": solveForTask_python}   
+    assert solver in solvers, "You must specify a valid solver. options are ocaml, pypy, or python." 
+
+    likelihoodModel = None
+    if solver == 'pypy' or solver == 'python':
+      # Use an all or nothing likelihood model.
+      likelihoodModel = AllOrNothingLikelihoodModel(timeout=evaluationTimeout) 
+      
+    solver = solvers[solver]
+
+    if not isinstance(g, dict):
+        g = {t: g for t in tasks}
+    task2grammar = g
+
+    # If we are not evaluating on held out testing tasks:
+    # Bin the tasks by request type and grammar
+    # If these are the same then we can enumerate for multiple tasks simultaneously
+    # If we are evaluating testing tasks:
+    # Make sure that each job corresponds to exactly one task
+    jobs = {}
+    for i, t in enumerate(tasks):
+        if testing:
+            k = (task2grammar[t], t.request, i)
+        else:
+            k = (task2grammar[t], t.request)
+        jobs[k] = jobs.get(k, []) + [t]
+
+    disableParallelism = len(jobs) == 1
+    parallelCallback = launchParallelProcess if not disableParallelism else lambda f, * \
+        a, **k: f(*a, **k)
+    if disableParallelism:
+        eprint("Disabling parallelism on the Python side because we only have one job.")
+        eprint("If you are using ocaml, there could still be parallelism.")
+
+    # Map from task to the shortest time to find a program solving it
+    bestSearchTime = {t: None for t in task2grammar}
+
+    lowerBounds = {k: 0. for k in jobs}
+
+    frontiers = {t: Frontier([], task=t) for t in task2grammar}
+
+    # For each job we keep track of how long we have been working on it
+    stopwatches = {t: Stopwatch() for t in jobs}
+
+    # Map from task to how many programs we enumerated for that task
+    taskToNumberOfPrograms = {t: 0 for t in tasks }
+
+    def numberOfHits(f):
+        return sum(e.logLikelihood > -0.01 for e in f)
+
+    def budgetIncrement(lb):
+        if True:
+            return 1.5
+        # Very heuristic - not sure what to do here
+        if lb < 24.:
+            return 1.
+        elif lb < 27.:
+            return 0.5
+        else:
+            return 0.25
+
+    def maximumFrontiers(j):
+        tasks = jobs[j]
+        return {t: maximumFrontier - numberOfHits(frontiers[t]) for t in tasks}
+
+    def allocateCPUs(n, tasks):
+        allocation = {t: 0 for t in tasks}
+        while n > 0:
+            for t in tasks:
+                # During testing we use exactly one CPU per task
+                if testing and allocation[t] > 0:
+                    return allocation
+                allocation[t] += 1
+                n -= 1
+                if n == 0:
+                    break
+        return allocation
+
+    def refreshJobs():
+        for k in list(jobs.keys()):
+            v = [t for t in jobs[k]
+                 if numberOfHits(frontiers[t]) < maximumFrontier
+                 and stopwatches[k].elapsed <= enumerationTimeout]
+            if v:
+                jobs[k] = v
+            else:
+                del jobs[k]
+
+    # Workers put their messages in here
+    q = Queue()
+
+    # How many CPUs are we using?
+    activeCPUs = 0
+
+    # How many CPUs was each job allocated?
+    id2CPUs = {}
+    # What job was each ID working on?
+    id2job = {}
+    nextID = 0
+
+    while True:
+        refreshJobs()
+        # Don't launch a job that we are already working on
+        # We run the stopwatch whenever the job is being worked on
+        # freeJobs are things that we are not working on but could be
+        freeJobs = [j for j in jobs if not stopwatches[j].running
+                    and stopwatches[j].elapsed < enumerationTimeout - 0.5]
+        if freeJobs and activeCPUs < CPUs:
+            # Allocate a CPU to each of the jobs that we have made the least
+            # progress on
+            freeJobs.sort(key=lambda j: lowerBounds[j])
+            # Launch some more jobs until all of the CPUs are being used
+            availableCPUs = CPUs - activeCPUs
+            allocation = allocateCPUs(availableCPUs, freeJobs)
+            for j in freeJobs:
+                if allocation[j] == 0:
+                    continue
+                g, request = j[:2]
+                bi = budgetIncrement(lowerBounds[j])
+                thisTimeout = enumerationTimeout - stopwatches[j].elapsed
+                eprint("(python) Launching %s (%d tasks) w/ %d CPUs. %f <= MDL < %f. Timeout %f." %
+                       (request, len(jobs[j]), allocation[j], lowerBounds[j], lowerBounds[j] + bi, thisTimeout))
+                stopwatches[j].start()
+                parallelCallback(wrapInThread(solver),
+                                 q=q, g=g, ID=nextID,
+                                 elapsedTime=stopwatches[j].elapsed,
+                                 CPUs=allocation[j],
+                                 tasks=jobs[j],
+                                 lowerBound=lowerBounds[j],
+                                 upperBound=lowerBounds[j] + bi,
+                                 budgetIncrement=bi,
+                                 timeout=thisTimeout,
+                                 evaluationTimeout=evaluationTimeout,
+                                 maximumFrontiers=maximumFrontiers(j),
+                                 testing=testing,
+                                 likelihoodModel=likelihoodModel)
+                id2CPUs[nextID] = allocation[j]
+                id2job[nextID] = j
+                nextID += 1
+
+                activeCPUs += allocation[j]
+                lowerBounds[j] += bi
+
+        # If nothing is running, and we just tried to launch jobs,
+        # then that means we are finished
+        if all(not s.running for s in stopwatches.values()):
+            break
+
+        # Wait to get a response
+        message = Bunch(dill.loads(q.get()))
+
+        if message.result == "failure":
+            eprint("PANIC! Exception in child worker:", message.exception)
+            eprint(message.stacktrace)
+            assert False
+        elif message.result == "success":
+            # Mark the CPUs is no longer being used and pause the stopwatch
+            activeCPUs -= id2CPUs[message.ID]
+            stopwatches[id2job[message.ID]].stop()
+
+            newFrontiers, searchTimes, pc = message.value
+            for t, f in newFrontiers.items():
+                oldBest = None if len(
+                    frontiers[t]) == 0 else frontiers[t].bestPosterior
+                frontiers[t] = frontiers[t].combine(f)
+                newBest = None if len(
+                    frontiers[t]) == 0 else frontiers[t].bestPosterior
+
+                taskToNumberOfPrograms[t] += pc
+
+                dt = searchTimes[t]
+                if dt is not None:
+                    if bestSearchTime[t] is None:
+                        bestSearchTime[t] = dt
+                    else:
+                        # newBest & oldBest should both be defined
+                        assert oldBest is not None
+                        assert newBest is not None
+                        newScore = newBest.logPrior + newBest.logLikelihood
+                        oldScore = oldBest.logPrior + oldBest.logLikelihood
+
+                        if newScore > oldScore:
+                            bestSearchTime[t] = dt
+                        elif newScore == oldScore:
+                            bestSearchTime[t] = min(bestSearchTime[t], dt)
+        else:
+            eprint("Unknown message result:", message.result)
+            assert False
+
+    eprint("We enumerated this many programs, for each task:\n\t",
+           list(taskToNumberOfPrograms.values()))
+
+    return [frontiers[t] for t in tasks], bestSearchTime
+
+def wrapInThread(f):
+    """
+    Returns a function that is designed to be run in a thread/threadlike process.
+    Result will be either put into the q
+    """
+    import dill
+
+    def _f(*a, **k):
+        q = k.pop("q")
+        ID = k.pop("ID")
+
+        try:
+            r = f(*a, **k)
+            q.put(dill.dumps({"result": "success",
+                   "ID": ID,
+                   "value": r}))
+        except Exception as e:
+            q.put(dill.dumps({"result": "failure",
+                   "exception": e,
+                   "stacktrace": traceback.format_exc(),
+                   "ID": ID}))
+            return
+    return _f
+
+
+def solveForTask_ocaml(_=None,
+                       elapsedTime=0.,
+                       CPUs=1,
+                       g=None, tasks=None,
+                       lowerBound=None, upperBound=None, budgetIncrement=None,
+                       timeout=None,
+                       testing=None, # FIXME: unused
+                       likelihoodModel=None,
+                       evaluationTimeout=None, maximumFrontiers=None):
+
+    import json
+
+    def taskMessage(t):
+        m = {
+            "examples": [{"inputs": list(xs), "output": y} for xs, y in t.examples],
+            "name": t.name,
+            "request": t.request.json(),
+            "maximumFrontier": maximumFrontiers[t]}
+        if hasattr(t, "specialTask"):
+            special, extra = t.specialTask
+            m["specialTask"] = special
+            m["extras"] = extra
+        return m
+
+
+    message = {"DSL": g.json(),
+               "tasks": [taskMessage(t)
+                         for t in tasks],
+
+               "programTimeout": evaluationTimeout,
+               "nc": CPUs,
+               "timeout": timeout,
+               "lowerBound": lowerBound,
+               "upperBound": upperBound,
+               "budgetIncrement": budgetIncrement,
+               "verbose": False,
+               "shatter": 5 if len(tasks) == 1 and "turtle" in str(tasks[0].request) else 10}
+
+    if hasattr(tasks[0], 'maxParameters') and tasks[0].maxParameters is not None:
+        message["maxParameters"] = tasks[0].maxParameters
+
+    message = json.dumps(message)
+    # uncomment this if you want to save the messages being sent to the solver
+    
+
+    try:
+        solver_file = os.path.join(get_root_dir(), 'solver')
+        process = subprocess.Popen(solver_file,
+                                   stdin=subprocess.PIPE,
+                                   stdout=subprocess.PIPE)
+        response, error = process.communicate(bytes(message, encoding="utf-8"))
+        response = json.loads(response.decode("utf-8"))
+    except OSError as exc:
+        raise exc
+
+    except:
+        print("response:", response)
+        print("error:", error)
+        with open("message", "w") as f:
+            f.write(message)
+        print("message,", message)
+        assert False, "MAX RAISE"
+
+
+    pc = response.get("number_enumerated",0)  # TODO
+    frontiers = {}
+    searchTimes = {}
+    for t in tasks:
+        solutions = response[t.name]
+        frontier = Frontier([FrontierEntry(program=p,
+                                           logLikelihood=e["logLikelihood"],
+                                           logPrior=g.logLikelihood(t.request, p))
+                             for e in solutions
+                             for p in [Program.parse(e["program"])]],
+                            task=t)
+        frontiers[t] = frontier
+        if frontier.empty:
+            searchTimes[t] = None
+        # This is subtle:
+        # The search time we report is actually not be minimum time to find any solution
+        # Rather it is the time to find the MAP solution
+        # This is important for regression problems,
+        # where we might find something with a good prior but bad likelihood early on,
+        # and only later discovered the good high likelihood program
+        else:
+            searchTimes[t] = min(
+                (e["logLikelihood"] + e["logPrior"],
+                 e["time"]) for e in solutions)[1] + elapsedTime
+
+    return frontiers, searchTimes, pc
+
+def solveForTask_pypy(_=None,
+                      elapsedTime=0.,
+                      g=None, task=None,
+                      lowerBound=None, upperBound=None, budgetIncrement=None,
+                      timeout=None,
+                      likelihoodModel=None,
+                      evaluationTimeout=None, maximumFrontier=None, testing=False):
+    return callCompiled(enumerateForTasks,
+                        g, tasks, likelihoodModel,
+                        timeout=timeout,
+                        testing=testing,
+                        elapsedTime=elapsedTime,
+                        evaluationTimeout=evaluationTimeout,
+                        maximumFrontiers=maximumFrontiers,
+                        budgetIncrement=budgetIncrement,
+                        lowerBound=lowerBound, upperBound=upperBound)
+
+def solveForTask_python(_=None,
+                        elapsedTime=0.,
+                        g=None, tasks=None,
+                        lowerBound=None, upperBound=None, budgetIncrement=None,
+                        timeout=None,
+                        CPUs=1,
+                        likelihoodModel=None,
+                        evaluationTimeout=None, maximumFrontiers=None, testing=False):
+    return enumerateForTasks(g, tasks, likelihoodModel,
+                             timeout=timeout,
+                             testing=testing,
+                             elapsedTime=elapsedTime,
+                             evaluationTimeout=evaluationTimeout,
+                             maximumFrontiers=maximumFrontiers,
+                             budgetIncrement=budgetIncrement,
+                             lowerBound=lowerBound, upperBound=upperBound)
+
+
+class EnumerationTimeout(Exception):
+    pass
+
+def enumerateForTasks(g, tasks, likelihoodModel, _=None,
+                      verbose=False,
+                      timeout=None,
+                      elapsedTime=0.,
+                      CPUs=1,
+                      testing=False, #unused
+                      evaluationTimeout=None,
+                      lowerBound=0.,
+                      upperBound=100.,
+                      budgetIncrement=1.0, maximumFrontiers=None):
+    assert timeout is not None, \
+        "enumerateForTasks: You must provide a timeout."
+
+    from time import time
+
+    request = tasks[0].request
+    assert all(t.request == request for t in tasks), \
+        "enumerateForTasks: Expected tasks to all have the same type"
+
+    maximumFrontiers = [maximumFrontiers[t] for t in tasks]
+    # store all of the hits in a priority queue
+    # we will never maintain maximumFrontier best solutions
+    hits = [PQ() for _ in tasks]
+
+    starting = time()
+    previousBudget = lowerBound
+    budget = lowerBound + budgetIncrement
+    try:
+        totalNumberOfPrograms = 0
+        while time() < starting + timeout and \
+                any(len(h) < mf for h, mf in zip(hits, maximumFrontiers)) and \
+                budget <= upperBound:
+            numberOfPrograms = 0
+
+            for prior, _, p in g.enumeration(Context.EMPTY, [], request,
+                                             maximumDepth=99,
+                                             upperBound=budget,
+                                             lowerBound=previousBudget):
+                descriptionLength = -prior
+                # Shouldn't see it on this iteration
+                assert descriptionLength <= budget
+                # Should already have seen it
+                assert descriptionLength > previousBudget
+
+                numberOfPrograms += 1
+                totalNumberOfPrograms += 1
+
+                for n in range(len(tasks)):
+                    task = tasks[n]
+
+                    #Warning:changed to max's new likelihood model situation
+                    #likelihood = task.logLikelihood(p, evaluationTimeout)
+                    #if invalid(likelihood):
+                        #continue
+                    success, likelihood = likelihoodModel.score(p, task)
+                    if not success:
+                        continue
+                        
+                    dt = time() - starting + elapsedTime
+                    priority = -(likelihood + prior)
+                    hits[n].push(priority,
+                                 (dt, FrontierEntry(program=p,
+                                                    logLikelihood=likelihood,
+                                                    logPrior=prior)))
+                    if len(hits[n]) > maximumFrontiers[n]:
+                        hits[n].popMaximum()
+
+                if timeout is not None and time() - starting > timeout:
+                    raise EnumerationTimeout
+
+            previousBudget = budget
+            budget += budgetIncrement
+
+            if budget > upperBound:
+                break
+    except EnumerationTimeout:
+        pass
+    frontiers = {tasks[n]: Frontier([e for _, e in hits[n]],
+                                    task=tasks[n])
+                 for n in range(len(tasks))}
+    searchTimes = {
+        tasks[n]: None if len(hits[n]) == 0 else \
+        min(t for t,_ in hits[n]) for n in range(len(tasks))}
+
+    return frontiers, searchTimes, totalNumberOfPrograms
+
+
+
+
+

dreamcoder/fragmentGrammar.py

@@ -0,0 +1,430 @@
+from dreamcoder.fragmentUtilities import *
+from dreamcoder.grammar import *
+from dreamcoder.program import *
+
+from itertools import chain
+import time
+
+
+class FragmentGrammar(object):
+    def __init__(self, logVariable, productions):
+        self.logVariable = logVariable
+        self.productions = productions
+        self.likelihoodCache = {}
+
+    def clearCache(self):
+        self.likelihoodCache = {}
+
+    def __repr__(self):
+        return "FragmentGrammar(logVariable={self.logVariable}, productions={self.productions}".format(
+            self=self)
+
+    def __str__(self):
+        def productionKey(xxx_todo_changeme):
+            (l, t, p) = xxx_todo_changeme
+            return not isinstance(p, Primitive), -l
+        return "\n".join(["%f\tt0\t$_" % self.logVariable] + ["%f\t%s\t%s" % (l, t, p)
+                                                              for l, t, p in sorted(self.productions, key=productionKey)])
+
+    def buildCandidates(self, context, environment, request):
+        candidates = []
+        variableCandidates = []
+        for l, t, p in self.productions:
+            try:
+                newContext, t = t.instantiate(context)
+                newContext = newContext.unify(t.returns(), request)
+                candidates.append((l, newContext,
+                                   t.apply(newContext),
+                                   p))
+            except UnificationFailure:
+                continue
+        for j, t in enumerate(environment):
+            try:
+                newContext = context.unify(t.returns(), request)
+                variableCandidates.append((newContext,
+                                           t.apply(newContext),
+                                           Index(j)))
+            except UnificationFailure:
+                continue
+        if variableCandidates:
+            z = math.log(len(variableCandidates))
+            for newContext, newType, index in variableCandidates:
+                candidates.append(
+                    (self.logVariable - z, newContext, newType, index))
+
+        z = lse([candidate[0] for candidate in candidates])
+        return [(l - z, c, t, p) for l, c, t, p in candidates]
+
+    def logLikelihood(self, request, expression):
+        _, l, _ = self._logLikelihood(Context.EMPTY, [], request, expression)
+        if invalid(l):
+            f = 'failures/likelihoodFailure%s.pickle' % (time() + getPID())
+            eprint("PANIC: Invalid log likelihood. expression:",
+                   expression, "tp:", request, "Exported to:", f)
+            with open(f, 'wb') as handle:
+                pickle.dump((self, request, expression), handle)
+            assert False
+        return l
+
+    def closedUses(self, request, expression):
+        _, l, u = self._logLikelihood(Context.EMPTY, [], request, expression)
+        return l, u
+
+    def _logLikelihood(self, context, environment, request, expression):
+        '''returns (context, log likelihood, uses)'''
+
+        # We can cash likelihood calculations faster whenever they don't involve type inference
+        # This is because they are guaranteed to not modify the context,
+        polymorphic = request.isPolymorphic or any(
+            v.isPolymorphic for v in environment)
+        # For some reason polymorphic caching slows it down
+        shouldDoCaching = not polymorphic
+
+        # Caching
+        if shouldDoCaching:
+            if polymorphic:
+                inTypes = canonicalTypes(
+                    [request.apply(context)] + [v.apply(context) for v in environment])
+            else:
+                inTypes = canonicalTypes([request] + environment)
+            cacheKey = (tuple(inTypes), expression)
+            if cacheKey in self.likelihoodCache:
+                outTypes, l, u = self.likelihoodCache[cacheKey]
+                context, instantiatedTypes = instantiateTypes(
+                    context, outTypes)
+                outRequest = instantiatedTypes[0]
+                outEnvironment = instantiatedTypes[1:]
+                # eprint("request:", request.apply(context), "environment:",
+                #        [ v.apply(context) for v in environment ])
+                # eprint("will be unified with: out request:",outRequest,"out environment",outEnvironment)
+                if polymorphic:
+                    context = context.unify(request, outRequest)
+                    for v, vp in zip(environment, outEnvironment):
+                        context = context.unify(v, vp)
+                return context, l, u
+
+        if request.isArrow():
+            if not isinstance(expression, Abstraction):
+                return (context, NEGATIVEINFINITY, Uses.empty)
+            return self._logLikelihood(context,
+                                       [request.arguments[0]] + environment,
+                                       request.arguments[1],
+                                       expression.body)
+
+        # Not a function type
+
+        # Construct and normalize the candidate productions
+        candidates = self.buildCandidates(context, environment, request)
+
+        # Consider each way of breaking the expression up into a
+        # function and arguments
+        totalLikelihood = NEGATIVEINFINITY
+        weightedUses = []
+
+        possibleVariables = float(int(any(isinstance(candidate, Index)
+                                          for _, _, _, candidate in candidates)))
+        possibleUses = {candidate: 1. for _, _, _, candidate in candidates
+                        if not isinstance(candidate, Index)}
+
+        for f, xs in expression.applicationParses():
+            for candidateLikelihood, newContext, tp, production in candidates:
+                variableBindings = {}
+                # This is a variable in the environment
+                if production.isIndex:
+                    if production != f:
+                        continue
+                else:
+                    try:
+                        newContext, fragmentType, variableBindings = \
+                            Matcher.match(newContext, production, f, len(xs))
+                        # This is necessary because the types of the variable
+                        # bindings and holes need to match up w/ request
+                        fragmentTypeTemplate = request
+                        for _ in xs:
+                            newContext, newVariable = newContext.makeVariable()
+                            fragmentTypeTemplate = arrow(
+                                newVariable, fragmentTypeTemplate)
+                        newContext = newContext.unify(
+                            fragmentType, fragmentTypeTemplate)
+                        # update the unified type
+                        tp = fragmentType.apply(newContext)
+                    except MatchFailure:
+                        continue
+
+                argumentTypes = tp.functionArguments()
+                if len(xs) != len(argumentTypes):
+                    # I think that this is some kind of bug. But I can't figure it out right now.
+                    # As a hack, count this as though it were a failure
+                    continue
+                    #raise GrammarFailure('len(xs) != len(argumentTypes): tp={}, xs={}'.format(tp, xs))
+
+                thisLikelihood = candidateLikelihood
+                if isinstance(production, Index):
+                    theseUses = Uses(possibleVariables=possibleVariables,
+                                     actualVariables=1.,
+                                     possibleUses=possibleUses.copy(),
+                                     actualUses={})
+                else:
+                    theseUses = Uses(possibleVariables=possibleVariables,
+                                     actualVariables=0.,
+                                     possibleUses=possibleUses.copy(),
+                                     actualUses={production: 1.})
+
+                # Accumulate likelihood from free variables and holes and
+                # arguments
+                for freeType, freeExpression in chain(
+                        variableBindings.values(), zip(argumentTypes, xs)):
+                    freeType = freeType.apply(newContext)
+                    newContext, expressionLikelihood, newUses = self._logLikelihood(
+                        newContext, environment, freeType, freeExpression)
+                    if expressionLikelihood is NEGATIVEINFINITY:
+                        thisLikelihood = NEGATIVEINFINITY
+                        break
+
+                    thisLikelihood += expressionLikelihood
+                    theseUses += newUses
+
+                if thisLikelihood is NEGATIVEINFINITY:
+                    continue
+
+                weightedUses.append((thisLikelihood, theseUses))
+                totalLikelihood = lse(totalLikelihood, thisLikelihood)
+
+                # Any of these new context objects should be equally good
+                context = newContext
+
+        if totalLikelihood is NEGATIVEINFINITY:
+            return context, totalLikelihood, Uses.empty
+        assert weightedUses != []
+
+        allUses = Uses.join(totalLikelihood, *weightedUses)
+
+        # memoize result
+        if shouldDoCaching:
+            outTypes = [request.apply(context)] + \
+                [v.apply(context) for v in environment]
+            outTypes = canonicalTypes(outTypes)
+            self.likelihoodCache[cacheKey] = (
+                outTypes, totalLikelihood, allUses)
+
+        return context, totalLikelihood, allUses
+
+    def expectedUses(self, frontiers):
+        if len(list(frontiers)) == 0:
+            return Uses()
+        likelihoods = [[(l + entry.logLikelihood, u)
+                        for entry in frontier
+                        for l, u in [self.closedUses(frontier.task.request, entry.program)]]
+                       for frontier in frontiers]
+        zs = (lse([l for l, _ in ls]) for ls in likelihoods)
+        return sum(math.exp(l - z) * u
+                   for z, frontier in zip(zs, likelihoods)
+                   for l, u in frontier)
+
+    def insideOutside(self, frontiers, pseudoCounts):
+        uses = self.expectedUses(frontiers)
+        return FragmentGrammar(log(uses.actualVariables +
+                                   pseudoCounts) -
+                               log(max(uses.possibleVariables, 1.)), [(log(uses.actualUses.get(p, 0.) +
+                                                                           pseudoCounts) -
+                                                                       log(uses.possibleUses.get(p, 0.) +
+                                                                           pseudoCounts), t, p) for _, t, p in self.productions])
+
+    def jointFrontiersLikelihood(self, frontiers):
+        return sum(lse([entry.logLikelihood + self.logLikelihood(frontier.task.request, entry.program)
+                        for entry in frontier])
+                   for frontier in frontiers)
+
+    def jointFrontiersMDL(self, frontiers, CPUs=1):
+        return sum(
+            parallelMap(
+                CPUs,
+                lambda frontier: max(
+                    entry.logLikelihood +
+                    self.logLikelihood(
+                        frontier.task.request,
+                        entry.program) for entry in frontier),
+                frontiers))
+
+    def __len__(self): return len(self.productions)
+
+    @staticmethod
+    def fromGrammar(g):
+        return FragmentGrammar(g.logVariable, g.productions)
+
+    def toGrammar(self):
+        return Grammar(self.logVariable, [(l, q.infer(), q)
+                                          for l, t, p in self.productions
+                                          for q in [defragment(p)]])
+
+    @property
+    def primitives(self): return [p for _, _, p in self.productions]
+
+    @staticmethod
+    def uniform(productions):
+        return FragmentGrammar(0., [(0., p.infer(), p) for p in productions])
+
+    def normalize(self):
+        z = lse([l for l, t, p in self.productions] + [self.logVariable])
+        return FragmentGrammar(self.logVariable - z,
+                               [(l - z, t, p) for l, t, p in self.productions])
+
+    def makeUniform(self):
+        return FragmentGrammar(0., [(0., p.infer(), p)
+                                    for _, _, p in self.productions])
+
+    def rescoreFrontier(self, frontier):
+        return Frontier([FrontierEntry(e.program,
+                                       logPrior=self.logLikelihood(frontier.task.request, e.program),
+                                       logLikelihood=e.logLikelihood)
+                         for e in frontier],
+                        frontier.task)
+
+    @staticmethod
+    def induceFromFrontiers(
+            g0,
+            frontiers,
+            _=None,
+            topK=1,
+            topk_use_only_likelihood=False,
+            pseudoCounts=1.0,
+            aic=1.0,
+            structurePenalty=0.001,
+            a=0,
+            CPUs=1):
+        _ = topk_use_only_likelihood # not used in python compressor
+        originalFrontiers = frontiers
+        frontiers = [frontier for frontier in frontiers if not frontier.empty]
+        eprint("Inducing a grammar from", len(frontiers), "frontiers")
+
+        bestGrammar = FragmentGrammar.fromGrammar(g0)
+        oldJoint = bestGrammar.jointFrontiersMDL(frontiers, CPUs=1)
+
+        # "restricted frontiers" only contain the top K according to the best grammar
+        def restrictFrontiers():
+            return parallelMap(
+                CPUs,
+                lambda f: bestGrammar.rescoreFrontier(f).topK(topK),
+                frontiers)
+        restrictedFrontiers = []
+
+        def grammarScore(g):
+            g = g.makeUniform().insideOutside(restrictedFrontiers, pseudoCounts)
+            likelihood = g.jointFrontiersMDL(restrictedFrontiers)
+            structure = sum(primitiveSize(p) for p in g.primitives)
+            score = likelihood - aic * len(g) - structurePenalty * structure
+            g.clearCache()
+            if invalid(score):
+                # FIXME: This should never occur but it does anyway
+                score = float('-inf')
+            return score, g
+
+        if aic is not POSITIVEINFINITY:
+            restrictedFrontiers = restrictFrontiers()
+            bestScore, _ = grammarScore(bestGrammar)
+            eprint("Starting score", bestScore)
+            while True:
+                restrictedFrontiers = restrictFrontiers()
+                fragments = [f
+                             for f in proposeFragmentsFromFrontiers(restrictedFrontiers, a, CPUs=CPUs)
+                             if f not in bestGrammar.primitives
+                             and defragment(f) not in bestGrammar.primitives]
+                eprint("Proposed %d fragments." % len(fragments))
+
+                candidateGrammars = [
+                    FragmentGrammar.uniform(
+                        bestGrammar.primitives +
+                        [fragment]) for fragment in fragments]
+                if not candidateGrammars:
+                    break
+
+                scoredFragments = parallelMap(CPUs, grammarScore, candidateGrammars,
+                                              # Each process handles up to 100
+                                              # grammars at a time, a "job"
+                                              chunksize=max(
+                                                  1, min(len(candidateGrammars) // CPUs, 100)),
+                                              # maxTasks: Maximum number of jobs allocated to a process
+                                              # This means that after evaluating this*chunk many grammars,
+                                              # we killed the process, freeing up its memory.
+                                              # In exchange we pay the cost of spawning a new process.
+                                              # We should play with this number,
+                                              # figuring out how big we can make it without
+                                              # running out of memory.
+                                              maxtasksperchild=5)
+                newScore, newGrammar = max(scoredFragments, key=lambda sg: sg[0])
+
+                if newScore <= bestScore:
+                    break
+                dS = newScore - bestScore
+                bestScore, bestGrammar = newScore, newGrammar
+                newPrimitiveLikelihood, newType, newPrimitive = bestGrammar.productions[-1]
+                expectedUses = bestGrammar.expectedUses(
+                    restrictedFrontiers).actualUses.get(newPrimitive, 0)
+                eprint(
+                    "New primitive of type %s\t%s\t\n(score = %f; dScore = %f; <uses> = %f)" %
+                    (newType, newPrimitive, newScore, dS, expectedUses))
+
+                # Rewrite the frontiers in terms of the new fragment
+                concretePrimitive = defragment(newPrimitive)
+                bestGrammar.productions[-1] = (newPrimitiveLikelihood,
+                                               concretePrimitive.tp,
+                                               concretePrimitive)
+                frontiers = parallelMap(
+                    CPUs, lambda frontier: bestGrammar.rescoreFrontier(
+                        RewriteFragments.rewriteFrontier(
+                            frontier, newPrimitive)), frontiers)
+                eprint(
+                    "\t(<uses> in rewritten frontiers: %f)" %
+                    (bestGrammar.expectedUses(frontiers).actualUses[concretePrimitive]))
+        else:
+            eprint("Skipping fragment proposals")
+
+        if False:
+            # Reestimate the parameters using the entire frontiers
+            bestGrammar = bestGrammar.makeUniform().insideOutside(frontiers, pseudoCounts)
+        elif True:
+            # Reestimate the parameters using the best programs
+            restrictedFrontiers = restrictFrontiers()
+            bestGrammar = bestGrammar.makeUniform().insideOutside(
+                restrictedFrontiers, pseudoCounts)
+        else:
+            # Use parameters that were found during search
+            pass
+
+        eprint("Old joint = %f\tNew joint = %f\n" %
+               (oldJoint, bestGrammar.jointFrontiersMDL(frontiers, CPUs=CPUs)))
+        # Return all of the frontiers, which have now been rewritten to use the
+        # new fragments
+        frontiers = {f.task: f for f in frontiers}
+        frontiers = [frontiers.get(f.task, f)
+                     for f in originalFrontiers]
+
+        productionUses = bestGrammar.expectedUses(
+            [f for f in frontiers if not f.empty]).actualUses
+        productionUses = {
+            p: productionUses.get(
+                p, 0.) for p in bestGrammar.primitives}
+        possibleUses = bestGrammar.expectedUses(
+            [f for f in frontiers if not f.empty]).possibleUses
+        possibleUses = {
+            p: possibleUses.get(
+                p, 0.) for p in bestGrammar.primitives}
+
+        for p in bestGrammar.primitives:
+            eprint("%f / %f\t%s" % (productionUses[p],
+                                    possibleUses[p],
+                                    p))
+
+        bestGrammar.clearCache()
+
+        grammar = bestGrammar.toGrammar()
+
+        if False and \
+           any(productionUses.get(p, 0) < 0.5 for p in grammar.primitives if p.isInvented):
+            uselessProductions = [ p for p in grammar.primitives                 
+                                   if p.isInvented and productionUses.get(p, 0) < 0.5]
+            eprint("The following invented primitives are no longer needed, removing them...")
+            eprint("\t" + "\t\n".join(map(str, uselessProductions)))
+            grammar = grammar.removeProductions(uselessProductions)
+
+        return grammar, frontiers

dreamcoder/fragmentUtilities.py

@@ -0,0 +1,405 @@
+from dreamcoder.type import *
+from dreamcoder.program import *
+from dreamcoder.frontier import *
+
+from collections import Counter
+
+
+class MatchFailure(Exception):
+    pass
+
+
+class Matcher(object):
+    def __init__(self, context):
+        self.context = context
+        self.variableBindings = {}
+
+    @staticmethod
+    def match(context, fragment, expression, numberOfArguments):
+        if not mightMatch(fragment, expression):
+            raise MatchFailure()
+        m = Matcher(context)
+        tp = fragment.visit(m, expression, [], numberOfArguments)
+        return m.context, tp, m.variableBindings
+
+    def application(
+            self,
+            fragment,
+            expression,
+            environment,
+            numberOfArguments):
+        '''returns tp of fragment.'''
+        if not isinstance(expression, Application):
+            raise MatchFailure()
+
+        ft = fragment.f.visit(
+            self,
+            expression.f,
+            environment,
+            numberOfArguments + 1)
+        xt = fragment.x.visit(self, expression.x, environment, 0)
+
+        self.context, returnType = self.context.makeVariable()
+        try:
+            self.context = self.context.unify(ft, arrow(xt, returnType))
+        except UnificationFailure:
+            raise MatchFailure()
+
+        return returnType.apply(self.context)
+
+    def index(self, fragment, expression, environment, numberOfArguments):
+        # This is a bound variable
+        surroundingAbstractions = len(environment)
+        if fragment.bound(surroundingAbstractions):
+            if expression == fragment:
+                return environment[fragment.i].apply(self.context)
+            else:
+                raise MatchFailure()
+
+        # This is a free variable
+        i = fragment.i - surroundingAbstractions
+
+        # Make sure that it doesn't refer to anything bound by a
+        # lambda in the fragment. Otherwise it cannot be safely lifted
+        # out of the fragment and preserve semantics
+        for fv in expression.freeVariables():
+            if fv < len(environment):
+                raise MatchFailure()
+
+        # The value is going to be lifted out of the fragment
+        try:
+            expression = expression.shift(-surroundingAbstractions)
+        except ShiftFailure:
+            raise MatchFailure()
+
+        # Wrap it in the appropriate number of lambda expressions & applications
+        # This is because everything has to be in eta-longform
+        if numberOfArguments > 0:
+            expression = expression.shift(numberOfArguments)
+            for j in reversed(range(numberOfArguments)):
+                expression = Application(expression, Index(j))
+            for _ in range(numberOfArguments):
+                expression = Abstraction(expression)
+
+        # Added to the bindings
+        if i in self.variableBindings:
+            (tp, binding) = self.variableBindings[i]
+            if binding != expression:
+                raise MatchFailure()
+        else:
+            self.context, tp = self.context.makeVariable()
+            self.variableBindings[i] = (tp, expression)
+        return tp
+
+    def abstraction(
+            self,
+            fragment,
+            expression,
+            environment,
+            numberOfArguments):
+        if not isinstance(expression, Abstraction):
+            raise MatchFailure()
+
+        self.context, argumentType = self.context.makeVariable()
+        returnType = fragment.body.visit(
+            self, expression.body, [argumentType] + environment, 0)
+
+        return arrow(argumentType, returnType)
+
+    def primitive(self, fragment, expression, environment, numberOfArguments):
+        if fragment != expression:
+            raise MatchFailure()
+        self.context, tp = fragment.tp.instantiate(self.context)
+        return tp
+
+    def invented(self, fragment, expression, environment, numberOfArguments):
+        if fragment != expression:
+            raise MatchFailure()
+        self.context, tp = fragment.tp.instantiate(self.context)
+        return tp
+
+    def fragmentVariable(
+            self,
+            fragment,
+            expression,
+            environment,
+            numberOfArguments):
+        raise Exception(
+            'Deprecated: matching against fragment variables. Convert fragment to canonical form to get rid of fragment variables.')
+
+
+def mightMatch(f, e, d=0):
+    '''Checks whether fragment f might be able to match against expression e'''
+    if f.isIndex:
+        if f.bound(d):
+            return f == e
+        return True
+    if f.isPrimitive or f.isInvented:
+        return f == e
+    if f.isAbstraction:
+        return e.isAbstraction and mightMatch(f.body, e.body, d + 1)
+    if f.isApplication:
+        return e.isApplication and mightMatch(
+            f.x, e.x, d) and mightMatch(
+            f.f, e.f, d)
+    assert False
+
+
+def canonicalFragment(expression):
+    '''
+    Puts a fragment into a canonical form:
+    1. removes all FragmentVariable's
+    2. renames all free variables based on depth first traversal
+    '''
+    return expression.visit(CanonicalVisitor(), 0)
+
+
+class CanonicalVisitor(object):
+    def __init__(self):
+        self.numberOfAbstractions = 0
+        self.mapping = {}
+
+    def fragmentVariable(self, e, d):
+        self.numberOfAbstractions += 1
+        return Index(self.numberOfAbstractions + d - 1)
+
+    def primitive(self, e, d): return e
+
+    def invented(self, e, d): return e
+
+    def application(self, e, d):
+        return Application(e.f.visit(self, d), e.x.visit(self, d))
+
+    def abstraction(self, e, d):
+        return Abstraction(e.body.visit(self, d + 1))
+
+    def index(self, e, d):
+        if e.bound(d):
+            return e
+        i = e.i - d
+        if i in self.mapping:
+            return Index(d + self.mapping[i])
+        self.mapping[i] = self.numberOfAbstractions
+        self.numberOfAbstractions += 1
+        return Index(self.numberOfAbstractions - 1 + d)
+
+
+def fragmentSize(f, boundVariableCost=0.1, freeVariableCost=0.01):
+    freeVariables = 0
+    leaves = 0
+    boundVariables = 0
+    for surroundingAbstractions, e in f.walk():
+        if isinstance(e, (Primitive, Invented)):
+            leaves += 1
+        if isinstance(e, Index):
+            if e.bound(surroundingAbstractions):
+                boundVariables += 1
+            else:
+                freeVariables += 1
+        assert not isinstance(e, FragmentVariable)
+    return leaves + boundVariableCost * \
+        boundVariables + freeVariableCost * freeVariables
+
+
+def primitiveSize(e):
+    if e.isInvented:
+        e = e.body
+    return fragmentSize(e)
+
+
+def defragment(expression):
+    '''Converts a fragment into an invented primitive'''
+    if isinstance(expression, (Primitive, Invented)):
+        return expression
+
+    expression = canonicalFragment(expression)
+
+    for _ in range(expression.numberOfFreeVariables):
+        expression = Abstraction(expression)
+
+    return Invented(expression)
+
+
+class RewriteFragments(object):
+    def __init__(self, fragment):
+        self.fragment = fragment
+        self.concrete = defragment(fragment)
+
+    def tryRewrite(self, e, numberOfArguments):
+        try:
+            context, t, bindings = Matcher.match(
+                Context.EMPTY, self.fragment, e, numberOfArguments)
+        except MatchFailure:
+            return None
+
+        assert frozenset(bindings.keys()) == frozenset(range(len(bindings))),\
+            "Perhaps the fragment is not in canonical form?"
+        e = self.concrete
+        for j in range(len(bindings) - 1, -1, -1):
+            _, b = bindings[j]
+            e = Application(e, b)
+        return e
+
+    def application(self, e, numberOfArguments):
+        e = Application(e.f.visit(self, numberOfArguments + 1),
+                        e.x.visit(self, 0))
+        return self.tryRewrite(e, numberOfArguments) or e
+
+    def index(self, e, numberOfArguments): return e
+
+    def invented(self, e, numberOfArguments): return e
+
+    def primitive(self, e, numberOfArguments): return e
+
+    def abstraction(self, e, numberOfArguments):
+        e = Abstraction(e.body.visit(self, 0))
+        return self.tryRewrite(e, numberOfArguments) or e
+
+    def rewrite(self, e): return e.visit(self, 0)
+
+    @staticmethod
+    def rewriteFrontier(frontier, fragment):
+        worker = RewriteFragments(fragment)
+        return Frontier([FrontierEntry(program=worker.rewrite(e.program),
+                                       logLikelihood=e.logLikelihood,
+                                       logPrior=e.logPrior,
+                                       logPosterior=e.logPosterior)
+                         for e in frontier],
+                        task=frontier.task)
+
+
+def proposeFragmentsFromFragment(f):
+    '''Abstracts out repeated structure within a single fragment'''
+    yield f
+    freeVariables = f.numberOfFreeVariables
+    closedSubtrees = Counter(
+        subtree for _,
+        subtree in f.walk() if not isinstance(
+            subtree,
+            Index) and subtree.closed)
+    del closedSubtrees[f]
+    for subtree, freq in closedSubtrees.items():
+        if freq < 2:
+            continue
+        yield canonicalFragment(f.substitute(subtree, Index(freeVariables)))
+
+
+def nontrivial(f):
+    if not isinstance(f, Application):
+        return False
+    # Curry
+    if isinstance(f.x, FragmentVariable):
+        return False
+    if isinstance(f.x, Index):
+        # Make sure that the index is used somewhere else
+        if not any(
+                isinstance(
+                    child,
+                    Index) and child.i -
+                surroundingAbstractions == f.x.i for surroundingAbstractions,
+                child in f.f.walk()):
+            return False
+
+    numberOfHoles = 0
+    numberOfVariables = 0
+    numberOfPrimitives = 0
+    for surroundingAbstractions, child in f.walk():
+        if isinstance(child, (Primitive, Invented)):
+            numberOfPrimitives += 1
+        if isinstance(child, FragmentVariable):
+            numberOfHoles += 1
+        if isinstance(child, Index) and child.free(surroundingAbstractions):
+            numberOfVariables += 1
+    #eprint("Fragment %s has %d calls and %d variables and %d primitives"%(f,numberOfHoles,numberOfVariables,numberOfPrimitives))
+
+    return numberOfPrimitives + 0.5 * \
+        (numberOfHoles + numberOfVariables) > 1.5 and numberOfPrimitives >= 1
+
+
+def violatesLaziness(fragment):
+    """
+    conditionals are lazy on the second and third arguments. this
+    invariant must be maintained by learned fragments.
+    """
+    for surroundingAbstractions, child in fragment.walkUncurried():
+        if not child.isApplication:
+            continue
+        f, xs = child.applicationParse()
+        if not (f.isPrimitive and f.name == "if"):
+            continue
+
+        # curried conditionals always violate laziness
+        if len(xs) != 3:
+            return True
+
+        # yes/no branches
+        y = xs[1]
+        n = xs[2]
+
+        return \
+            any(yc.isIndex and yc.i >= yd
+                for yd, yc in y.walk(surroundingAbstractions)) or \
+            any(nc.isIndex and nc.i >= nd
+                for nd, nc in n.walk(surroundingAbstractions))
+
+    return False
+
+
+def proposeFragmentsFromProgram(p, arity):
+
+    def fragment(expression, a, toplevel=True):
+        """Generates fragments that unify with expression"""
+
+        if a == 1:
+            yield FragmentVariable.single
+        if a == 0:
+            yield expression
+            return
+
+        if isinstance(expression, Abstraction):
+            # Symmetry breaking: (\x \y \z ... f(x,y,z,...)) defragments to be
+            # the same as f(x,y,z,...)
+            if not toplevel:
+                for b in fragment(expression.body, a, toplevel=False):
+                    yield Abstraction(b)
+        elif isinstance(expression, Application):
+            for fa in range(a + 1):
+                for f in fragment(expression.f, fa, toplevel=False):
+                    for x in fragment(expression.x, a - fa, toplevel=False):
+                        yield Application(f, x)
+        else:
+            assert isinstance(expression, (Invented, Primitive, Index))
+
+    def fragments(expression, a):
+        """Generates fragments that unify with subexpressions of expression"""
+
+        yield from fragment(expression, a)
+        if isinstance(expression, Application):
+            curry = True
+            if curry:
+                yield from fragments(expression.f, a)
+                yield from fragments(expression.x, a)
+            else:
+                # Pretend that it is not curried
+                function, arguments = expression.applicationParse()
+                yield from fragments(function, a)
+                for argument in arguments:
+                    yield from fragments(argument, a)
+        elif isinstance(expression, Abstraction):
+            yield from fragments(expression.body, a)
+        else:
+            assert isinstance(expression, (Invented, Primitive, Index))
+
+    return {canonicalFragment(f) for b in range(arity + 1)
+            for f in fragments(p, b) if nontrivial(f)}
+
+
+def proposeFragmentsFromFrontiers(frontiers, a, CPUs=1):
+    fragmentsFromEachFrontier = parallelMap(
+        CPUs, lambda frontier: {
+            fp for entry in frontier.entries for f in proposeFragmentsFromProgram(
+                entry.program, a) for fp in proposeFragmentsFromFragment(f)}, frontiers)
+    allFragments = Counter(f for frontierFragments in fragmentsFromEachFrontier
+                           for f in frontierFragments)
+    return [fragment for fragment, frequency in allFragments.items()
+            if frequency >= 2 and fragment.wellTyped() and nontrivial(fragment)]

dreamcoder/frontier.py

@@ -0,0 +1,247 @@
+from dreamcoder.utilities import *
+from dreamcoder.program import *
+from dreamcoder.task import Task
+
+
+class FrontierEntry(object):
+    def __init__(
+            self,
+            program,
+            _=None,
+            logPrior=None,
+            logLikelihood=None,
+            logPosterior=None):
+        self.logPosterior = logPrior + logLikelihood if logPosterior is None else logPosterior
+        self.program = program
+        self.logPrior = logPrior
+        self.logLikelihood = logLikelihood
+
+    def __repr__(self):
+        return "FrontierEntry(program={self.program}, logPrior={self.logPrior}, logLikelihood={self.logLikelihood}".format(
+            self=self)
+
+    def strip_primitive_values(self):
+        return FrontierEntry(program=strip_primitive_values(self.program),
+                             logPrior=self.logPrior,
+                             logPosterior=self.logPosterior,
+                             logLikelihood=self.logLikelihood)
+    def unstrip_primitive_values(self):
+        return FrontierEntry(program=unstrip_primitive_values(self.program),
+                             logPrior=self.logPrior,
+                             logPosterior=self.logPosterior,
+                             logLikelihood=self.logLikelihood)
+
+
+class Frontier(object):
+    def __init__(self, frontier, task):
+        self.entries = frontier
+        self.task = task
+
+    def __repr__(
+        self): return "Frontier(entries={self.entries}, task={self.task})".format(self=self)
+
+    def __iter__(self): return iter(self.entries)
+
+    def __len__(self): return len(self.entries)
+
+    def json(self):
+        return {"request": self.task.request.json(),
+                "task": str(self.task),
+                "programs": [{"program": str(e.program),
+                              "logLikelihood": e.logLikelihood}
+                             for e in self ]}
+
+    def strip_primitive_values(self):
+        return Frontier([e.strip_primitive_values() for e in self.entries ],
+                        self.task)
+    def unstrip_primitive_values(self):
+        return Frontier([e.unstrip_primitive_values() for e in self.entries ],
+                        self.task)
+
+    DUMMYFRONTIERCOUNTER = 0
+
+    @staticmethod
+    def dummy(program, logLikelihood=0., logPrior=0., tp=None):
+        """Creates a dummy frontier containing just this program"""
+        if not tp:
+            tp = program.infer().negateVariables()
+
+        t = Task(
+            "<dummy %d: %s>" %
+            (Frontier.DUMMYFRONTIERCOUNTER,
+             str(program)),
+            tp,
+            [])
+        f = Frontier([FrontierEntry(program=program,
+                                    logLikelihood=logLikelihood,
+                                    logPrior=logPrior)],
+                     task=t)
+        Frontier.DUMMYFRONTIERCOUNTER += 1
+        return f
+
+    def marginalLikelihood(self):
+        return lse([e.logPrior + e.logLikelihood for e in self])
+
+    def temperature(self,T):
+        """Divides prior by T"""
+        return Frontier([ FrontierEntry(program=e.program,
+                                        logPrior=e.logPrior/T,
+                                        logLikelihood=e.logLikelihood)
+                          for e in self],
+                        task=self.task)
+                                        
+
+    def normalize(self):
+        z = self.marginalLikelihood()
+        newEntries = [
+            FrontierEntry(
+                program=e.program,
+                logPrior=e.logPrior,
+                logLikelihood=e.logLikelihood,
+                logPosterior=e.logPrior +
+                e.logLikelihood -
+                z) for e in self]
+        newEntries.sort(key=lambda e: e.logPosterior, reverse=True)
+        return Frontier(newEntries,
+                        self.task)
+
+    def expectedProductionUses(self, g):
+        """Returns a vector of the expected number of times each production was used"""
+        import numpy as np
+
+        this = g.rescoreFrontier(self).normalize()
+        ps = list(sorted(g.primitives, key=str))
+        features = np.zeros(len(ps))
+        
+        for j, p in enumerate(ps):
+            for e in this:
+                w = math.exp(e.logPosterior)
+                features[j] += w * sum(child == p
+                                       for _, child in e.program.walk() )
+            if not p.isInvented: features[j] *= 0.3
+        return features
+            
+
+    def removeZeroLikelihood(self):
+        self.entries = [
+            e for e in self.entries if e.logLikelihood != float('-inf')]
+        return self
+
+    def topK(self, k):
+        if k == 0: return Frontier([], self.task)
+        if k < 0: return self            
+        newEntries = sorted(self.entries,
+                            key=lambda e: (-e.logPosterior, str(e.program)))
+        return Frontier(newEntries[:k], self.task)
+
+    def sample(self):
+        """Samples an entry from a frontier"""
+        return sampleLogDistribution([(e.logLikelihood + e.logPrior, e)
+                                      for e in self])
+
+    @property
+    def bestPosterior(self):
+        return min(self.entries,
+                   key=lambda e: (-e.logPosterior, str(e.program)))
+
+    def replaceWithSupervised(self, g):
+        assert self.task.supervision is not None
+        return g.rescoreFrontier(Frontier([FrontierEntry(self.task.supervision,
+                                                         logLikelihood=0., logPrior=0.)],
+                                          task=self.task))
+
+    @property
+    def bestll(self):
+        best = max(self.entries,
+                   key=lambda e: e.logLikelihood)
+        return best.logLikelihood
+
+
+    @property
+    def empty(self): return self.entries == []
+
+    @staticmethod
+    def makeEmpty(task):
+        return Frontier([], task=task)
+
+    def summarize(self):
+        if self.empty:
+            return "MISS " + self.task.name
+        best = self.bestPosterior
+        return "HIT %s w/ %s ; log prior = %f ; log likelihood = %f" % (
+            self.task.name, best.program, best.logPrior, best.logLikelihood)
+
+    def summarizeFull(self):
+        if self.empty:
+            return "MISS " + self.task.name
+        return "\n".join([self.task.name] +
+                         ["%f\t%s" % (e.logPosterior, e.program)
+                          for e in self.normalize()])
+
+    @staticmethod
+    def describe(frontiers):
+        numberOfHits = sum(not f.empty for f in frontiers)
+        if numberOfHits > 0:
+            averageLikelihood = sum(
+                f.bestPosterior.logPrior for f in frontiers if not f.empty) / numberOfHits
+        else:
+            averageLikelihood = 0
+        return "\n".join([f.summarize() for f in frontiers] +
+                         ["Hits %d/%d tasks" % (numberOfHits, len(frontiers))] +
+                         ["Average description length of a program solving a task: %f nats" % (-averageLikelihood)])
+
+    def combine(self, other, tolerance=0.01):
+        '''Takes the union of the programs in each of the frontiers'''
+        assert self.task == other.task
+
+        foundDifference = False
+
+        x = {e.program: e for e in self}
+        y = {e.program: e for e in other}
+        programs = set(x.keys()) | set(y.keys())
+        union = []
+        for p in programs:
+            if p in x:
+                e1 = x[p]
+                if p in y:
+                    e2 = y[p]
+                    if abs(e1.logPrior - e2.logPrior) > tolerance:
+                        eprint(
+                            "WARNING: Log priors differed during frontier combining: %f vs %f" %
+                            (e1.logPrior, e2.logPrior))
+                        eprint("WARNING: \tThe program is", p)
+                        eprint()
+                    if abs(e1.logLikelihood - e2.logLikelihood) > tolerance:
+                        foundDifference = True
+                        eprint(
+                            "WARNING: Log likelihoods deferred for %s: %f & %f" %
+                            (p, e1.logLikelihood, e2.logLikelihood))
+                        if hasattr(self.task, 'BIC'):
+                            eprint("\t%d examples, BIC=%f, parameterPenalty=%f, n parameters=%d, correct likelihood=%f" %
+                                   (len(self.task.examples),
+                                    self.task.BIC,
+                                    self.task.BIC * math.log(len(self.task.examples)),
+                                    substringOccurrences("REAL", str(p)),
+                                    substringOccurrences("REAL", str(p)) * self.task.BIC * math.log(len(self.task.examples))))
+                            e1.logLikelihood = - \
+                                substringOccurrences("REAL", str(p)) * self.task.BIC * math.log(len(self.task.examples))
+                            e2.logLikelihood = e1.logLikelihood
+
+                        e1 = FrontierEntry(
+                            program=e1.program,
+                            logLikelihood=(
+                                e1.logLikelihood +
+                                e2.logLikelihood) /
+                            2,
+                            logPrior=e1.logPrior)
+            else:
+                e1 = y[p]
+            union.append(e1)
+
+        if foundDifference:
+            eprint(
+                "WARNING: Log likelihoods differed for the same program on the task %s.\n" %
+                (self.task.name),
+                "\tThis is acceptable only if the likelihood model is stochastic. Took the geometric mean of the likelihoods.")
+
+        return Frontier(union, self.task)

dreamcoder/grammar.py

@@ -0,0 +1,1308 @@
+from collections import defaultdict
+
+from dreamcoder.frontier import *
+from dreamcoder.program import *
+from dreamcoder.type import *
+from dreamcoder.utilities import *
+
+import time
+
+class GrammarFailure(Exception):
+    pass
+
+class SketchEnumerationFailure(Exception):
+    pass
+
+class NoCandidates(Exception):
+    pass
+
+
+class Grammar(object):
+    def __init__(self, logVariable, productions, continuationType=None):
+        self.logVariable = logVariable
+        self.productions = productions
+
+        self.continuationType = continuationType
+
+        self.expression2likelihood = dict((p, l) for l, _, p in productions)
+        self.expression2likelihood[Index(0)] = self.logVariable
+
+    def randomWeights(self, r):
+        """returns a new grammar with random weights drawn from r. calls `r` w/ old weight"""
+        return Grammar(logVariable=r(self.logVariable),
+                       productions=[(r(l),t,p)
+                                    for l,t,p in self.productions ],
+                       continuationType=self.continuationType)
+
+    def strip_primitive_values(self):
+        return Grammar(logVariable=self.logVariable,
+                       productions=[(l,t,strip_primitive_values(p))
+                                    for l,t,p in self.productions ],
+                       continuationType=self.continuationType)
+
+    def unstrip_primitive_values(self):
+        return Grammar(logVariable=self.logVariable,
+                       productions=[(l,t,unstrip_primitive_values(p))
+                                    for l,t,p in self.productions ],
+                       continuationType=self.continuationType)
+
+    def __setstate__(self, state):
+        """
+        Legacy support for loading grammar objects without the imperative type filled in
+        """
+        assert 'logVariable' in state
+        assert 'productions' in state
+        if 'continuationType' in state:
+            continuationType = state['continuationType']
+        else:
+            if any( 'turtle' in str(t) for l,t,p in state['productions'] ):
+                continuationType = baseType("turtle")
+            elif any( 'tower' in str(t) for l,t,p in state['productions'] ):
+                continuationType = baseType("tower")
+            else:
+                continuationType = None
+                
+        self.__init__(state['logVariable'], state['productions'], continuationType=continuationType)
+
+    @staticmethod
+    def fromProductions(productions, logVariable=0.0, continuationType=None):
+        """Make a grammar from primitives and their relative logpriors."""
+        return Grammar(logVariable, [(l, p.infer(), p)
+                                     for l, p in productions],
+                       continuationType=continuationType)
+
+    @staticmethod
+    def uniform(primitives, continuationType=None):
+        return Grammar(0.0, [(0.0, p.infer(), p) for p in primitives], continuationType=continuationType)
+
+    def __len__(self): return len(self.productions)
+
+    def __str__(self):
+        def productionKey(xxx_todo_changeme):
+            (l, t, p) = xxx_todo_changeme
+            return not isinstance(p, Primitive), l is not None and -l
+        if self.continuationType is not None:
+            lines = ["continuation : %s"%self.continuationType]
+        else:
+            lines = []
+        lines += ["%f\tt0\t$_" % self.logVariable]
+        for l, t, p in sorted(self.productions, key=productionKey):
+            if l is not None:
+                l = "%f\t%s\t%s" % (l, t, p)
+            else:
+                l = "-Inf\t%s\t%s" % (t, p)
+            if not t.isArrow() and isinstance(p, Invented):
+                try:
+                    l += "\teval = %s" % (p.evaluate([]))
+                except BaseException:
+                    pass
+
+            lines.append(l)
+        return "\n".join(lines)
+
+    def json(self):
+        j = {"logVariable": self.logVariable,
+             "productions": [{"expression": str(p), "logProbability": l}
+                             for l, _, p in self.productions]}
+        if self.continuationType is not None:
+            j["continuationType"] = self.continuationType.json()
+        return j
+
+    def _immutable_code(self): return self.logVariable, tuple(self.productions)
+
+    def __eq__(self, o): return self._immutable_code() == o._immutable_code()
+
+    def __ne__(self, o): return not (self == o)
+
+    def __hash__(self): return hash(self._immutable_code())
+
+    @property
+    def primitives(self):
+        return [p for _, _, p in self.productions]
+
+    def removeProductions(self, ps):
+        return Grammar(
+            self.logVariable, [
+                (l, t, p) for (
+                    l, t, p) in self.productions if p not in ps],
+            continuationType=self.continuationType)
+
+    def buildCandidates(self, request, context, environment,
+                        # Should the log probabilities be normalized?
+                        normalize=True,
+                        # Should be returned a table mapping primitives to
+                        # their candidate entry?
+                        returnTable=False,
+                        # Should we return probabilities vs log probabilities?
+                        returnProbabilities=False,
+                        # Must be a leaf (have no arguments)?
+                        mustBeLeaf=False):
+        """Primitives that are candidates for being used given a requested type
+        If returnTable is false (default): returns [((log)likelihood, tp, primitive, context)]
+        if returntable is true: returns {primitive: ((log)likelihood, tp, context)}"""
+        if returnProbabilities:
+            assert normalize
+
+        candidates = []
+        variableCandidates = []
+        for l, t, p in self.productions:
+            try:
+                newContext, t = t.instantiate(context)
+                newContext = newContext.unify(t.returns(), request)
+                t = t.apply(newContext)
+                if mustBeLeaf and t.isArrow():
+                    continue
+                candidates.append((l, t, p, newContext))
+            except UnificationFailure:
+                continue
+        for j, t in enumerate(environment):
+            try:
+                newContext = context.unify(t.returns(), request)
+                t = t.apply(newContext)
+                if mustBeLeaf and t.isArrow():
+                    continue
+                variableCandidates.append((t, Index(j), newContext))
+            except UnificationFailure:
+                continue
+
+        if self.continuationType == request:
+            terminalIndices = [v.i for t,v,k in variableCandidates if not t.isArrow()]
+            if terminalIndices:
+                smallestIndex = Index(min(terminalIndices))
+                variableCandidates = [(t,v,k) for t,v,k in variableCandidates
+                                      if t.isArrow() or v == smallestIndex]
+            
+        candidates += [(self.logVariable - log(len(variableCandidates)), t, p, k)
+                       for t, p, k in variableCandidates]
+        if candidates == []:
+            raise NoCandidates()
+        #eprint("candidates inside buildCandidates before norm:")
+        #eprint(candidates)
+
+        if normalize:
+            z = lse([l for l, t, p, k in candidates])
+            if returnProbabilities:
+                candidates = [(exp(l - z), t, p, k)
+                              for l, t, p, k in candidates]
+            else:
+                candidates = [(l - z, t, p, k) for l, t, p, k in candidates]
+
+        #eprint("candidates inside buildCandidates after norm:")
+        #eprint(candidates)
+
+        if returnTable:
+            return {p: (l, t, k) for l, t, p, k in candidates}
+        else:
+            return candidates
+
+
+    def sample(self, request, maximumDepth=6, maxAttempts=None):
+        attempts = 0
+
+        while True:
+            try:
+                _, e = self._sample(
+                    request, Context.EMPTY, [], maximumDepth=maximumDepth)
+                return e
+            except NoCandidates:
+                if maxAttempts is not None:
+                    attempts += 1
+                    if attempts > maxAttempts:
+                        return None
+                continue
+
+    def _sample(self, request, context, environment, maximumDepth):
+        if request.isArrow():
+            context, expression = self._sample(
+                request.arguments[1], context, [
+                    request.arguments[0]] + environment, maximumDepth)
+            return context, Abstraction(expression)
+
+        candidates = self.buildCandidates(request, context, environment,
+                                          normalize=True,
+                                          returnProbabilities=True,
+                                          # Force it to terminate in a
+                                          # leaf; a primitive with no
+                                          # function arguments
+                                          mustBeLeaf=maximumDepth <= 1)
+        #eprint("candidates:")
+        #eprint(candidates)
+        newType, chosenPrimitive, context = sampleDistribution(candidates)
+
+        # Sample the arguments
+        xs = newType.functionArguments()
+        returnValue = chosenPrimitive
+
+        for x in xs:
+            x = x.apply(context)
+            context, x = self._sample(x, context, environment, maximumDepth - 1)
+            returnValue = Application(returnValue, x)
+
+        return context, returnValue
+
+    def likelihoodSummary(self, context, environment, request, expression, silent=False):
+        if request.isArrow():
+            if not isinstance(expression, Abstraction):
+                if not silent:
+                    eprint("Request is an arrow but I got", expression)
+                return context, None
+            return self.likelihoodSummary(context,
+                                          [request.arguments[0]] + environment,
+                                          request.arguments[1],
+                                          expression.body,
+                                          silent=silent)
+        # Build the candidates
+        candidates = self.buildCandidates(request, context, environment,
+                                          normalize=False,
+                                          returnTable=True)
+
+        # A list of everything that would have been possible to use here
+        possibles = [p for p in candidates.keys() if not p.isIndex]
+        numberOfVariables = sum(p.isIndex for p in candidates.keys())
+        if numberOfVariables > 0:
+            possibles += [Index(0)]
+
+        f, xs = expression.applicationParse()
+
+        if f not in candidates:
+            if self.continuationType is not None and f.isIndex:
+                ls = LikelihoodSummary()
+                ls.constant = NEGATIVEINFINITY
+                return ls
+            
+            if not silent:
+                eprint(f, "Not in candidates")
+                eprint("Candidates is", candidates)
+                #eprint("grammar:", grammar.productions)
+                eprint("request is", request)
+                eprint("xs", xs)
+                eprint("environment", environment)
+                assert False
+            return context, None
+
+        thisSummary = LikelihoodSummary()
+        thisSummary.record(f, possibles,
+                           constant= -math.log(numberOfVariables) if f.isIndex else 0)
+
+        _, tp, context = candidates[f]
+        argumentTypes = tp.functionArguments()
+        if len(xs) != len(argumentTypes):
+            eprint("PANIC: not enough arguments for the type")
+            eprint("request", request)
+            eprint("tp", tp)
+            eprint("expression", expression)
+            eprint("xs", xs)
+            eprint("argumentTypes", argumentTypes)
+            # This should absolutely never occur
+            raise GrammarFailure((context, environment, request, expression))
+
+        for argumentType, argument in zip(argumentTypes, xs):
+            argumentType = argumentType.apply(context)
+            context, newSummary = self.likelihoodSummary(
+                context, environment, argumentType, argument, silent=silent)
+            if newSummary is None:
+                return context, None
+            thisSummary.join(newSummary)
+
+        return context, thisSummary
+
+    def bestFirstEnumeration(self, request):
+        from heapq import heappush, heappop
+
+        pq = []
+
+        def choices(parentCost, xs):
+            for c, x in xs:
+                heappush(pq, (parentCost + c, x))
+
+        def g(parentCost, request, _=None,
+              context=None, environment=[],
+              k=None):
+            """
+            k is a continuation.
+            k: Expects to be called with MDL, context, expression.
+            """
+
+            assert k is not None
+            if context is None:
+                context = Context.EMPTY
+
+            if request.isArrow():
+                g(parentCost,
+                  request.arguments[1],
+                  context=context,
+                  environment=[request.arguments[0]] + environment,
+                    k=lambda MDL,
+                    newContext,
+                    p: k(MDL,
+                         newContext,
+                         Abstraction(p)))
+            else:
+                candidates = self.buildCandidates(request,
+                                                  context,
+                                                  environment,
+                                                  normalize=True,
+                                                  returnProbabilities=False,
+                                                  returnTable=True)
+                choices(parentCost,
+                        [(-f_ll_tp_newContext[1][0],
+                          lambda: ga(parentCost - f_ll_tp_newContext[1][0],
+                                     f_ll_tp_newContext[0],
+                                     f_ll_tp_newContext[1][1].functionArguments(),
+                                     context=f_ll_tp_newContext[1][2],
+                                     environment=environment,
+                                     k=k)) for f_ll_tp_newContext in iter(candidates.items())])
+
+        def ga(costSoFar, f, argumentTypes, _=None,
+               context=None, environment=None,
+               k=None):
+            if argumentTypes == []:
+                k(costSoFar, context, f)
+            else:
+                t1 = argumentTypes[0].apply(context)
+                g(costSoFar, t1, context=context, environment=environment,
+                  k=lambda newCost, newContext, argument:
+                  ga(newCost, Application(f, argument), argumentTypes[1:],
+                     context=newContext, environment=environment,
+                     k=k))
+
+        def receiveResult(MDL, _, expression):
+            heappush(pq, (MDL, expression))
+
+        g(0., request, context=Context.EMPTY, environment=[], k=receiveResult)
+        frontier = []
+        while len(frontier) < 10**3:
+            MDL, action = heappop(pq)
+            if isinstance(action, Program):
+                expression = action
+                frontier.append(expression)
+                #eprint("Enumerated program",expression,-MDL,self.closedLogLikelihood(request, expression))
+            else:
+                action()
+
+    def closedLikelihoodSummary(self, request, expression, silent=False):
+        try:
+            context, summary = self.likelihoodSummary(Context.EMPTY, [], request, expression, silent=silent)
+        except GrammarFailure as e:
+            failureExport = 'failures/grammarFailure%s.pickle' % (
+                time.time() + getPID())
+            eprint("PANIC: Grammar failure, exporting to ", failureExport)
+            with open(failureExport, 'wb') as handle:
+                pickle.dump((e, self, request, expression), handle)
+            assert False
+
+        return summary
+
+    def logLikelihood(self, request, expression):
+        summary = self.closedLikelihoodSummary(request, expression)
+        if summary is None:
+            eprint(
+                "FATAL: program [ %s ] does not have a likelihood summary." %
+                expression, "r = ", request, "\n", self)
+            assert False
+        return summary.logLikelihood(self)
+
+    def rescoreFrontier(self, frontier):
+        return Frontier([FrontierEntry(e.program,
+                                       logPrior=self.logLikelihood(frontier.task.request, e.program),
+                                       logLikelihood=e.logLikelihood)
+                         for e in frontier],
+                        frontier.task)
+
+    def productionUses(self, frontiers):
+        """Returns the expected number of times that each production was used. {production: expectedUses}"""
+        frontiers = [self.rescoreFrontier(f).normalize()
+                     for f in frontiers if not f.empty]
+        uses = {p: 0. for p in self.primitives}
+        for f in frontiers:
+            for e in f:
+                summary = self.closedLikelihoodSummary(f.task.request,
+                                                       e.program)
+                for p, u in summary.uses:
+                    uses[p] += u * math.exp(e.logPosterior)
+        return uses
+
+    def insideOutside(self, frontiers, pseudoCounts, iterations=1):
+        # Replace programs with (likelihood summary, uses)
+        frontiers = [ Frontier([ FrontierEntry((summary, summary.toUses()),
+                                               logPrior=summary.logLikelihood(self),
+                                               logLikelihood=e.logLikelihood)
+                                 for e in f
+                                 for summary in [self.closedLikelihoodSummary(f.task.request, e.program)] ],
+                               task=f.task)
+                      for f in frontiers ]
+
+        g = self
+        for i in range(iterations):
+            u = Uses()
+            for f in frontiers:
+                f = f.normalize()
+                for e in f:
+                    _, eu = e.program
+                    u += math.exp(e.logPosterior) * eu
+
+            lv = math.log(u.actualVariables + pseudoCounts) - \
+                 math.log(u.possibleVariables + pseudoCounts)
+            g = Grammar(lv,
+                        [ (math.log(u.actualUses.get(p,0.) + pseudoCounts) - \
+                           math.log(u.possibleUses.get(p,0.) + pseudoCounts),
+                           t,p)
+                          for _,t,p in g.productions ],
+                        continuationType=self.continuationType)
+            if i < iterations - 1:
+                frontiers = [Frontier([ FrontierEntry((summary, uses),
+                                                      logPrior=summary.logLikelihood(g),
+                                                      logLikelihood=e.logLikelihood)
+                                        for e in f
+                                        for (summary, uses) in [e.program] ],
+                                      task=f.task)
+                             for f in frontiers ]
+        return g
+
+    def frontierMDL(self, frontier):
+        return max( e.logLikelihood + self.logLikelihood(frontier.task.request, e.program)
+                    for e in frontier )                
+
+
+    def enumeration(self,context,environment,request,upperBound,
+                    maximumDepth=20,
+                    lowerBound=0.):
+        '''Enumerates all programs whose MDL satisfies: lowerBound <= MDL < upperBound'''
+        if upperBound < 0 or maximumDepth == 1:
+            return
+
+        if request.isArrow():
+            v = request.arguments[0]
+            for l, newContext, b in self.enumeration(context, [v] + environment,
+                                                     request.arguments[1],
+                                                     upperBound=upperBound,
+                                                     lowerBound=lowerBound,
+                                                     maximumDepth=maximumDepth):
+                yield l, newContext, Abstraction(b)
+
+        else:
+            candidates = self.buildCandidates(request, context, environment,
+                                              normalize=True)
+
+            for l, t, p, newContext in candidates:
+                mdl = -l
+                if not (mdl < upperBound):
+                    continue
+
+                xs = t.functionArguments()
+                for aL, aK, application in\
+                    self.enumerateApplication(newContext, environment, p, xs,
+                                              upperBound=upperBound + l,
+                                              lowerBound=lowerBound + l,
+                                              maximumDepth=maximumDepth - 1):
+                    yield aL + l, aK, application
+
+    def enumerateApplication(self, context, environment,
+                             function, argumentRequests,
+                             # Upper bound on the description length of all of
+                             # the arguments
+                             upperBound,
+                             # Lower bound on the description length of all of
+                             # the arguments
+                             lowerBound=0.,
+                             maximumDepth=20,
+                             originalFunction=None,
+                             argumentIndex=0):
+        if upperBound < 0. or maximumDepth == 1:
+            return
+        if originalFunction is None:
+            originalFunction = function
+
+        if argumentRequests == []:
+            if lowerBound <= 0. and 0. < upperBound:
+                yield 0., context, function
+            else:
+                return
+        else:
+            argRequest = argumentRequests[0].apply(context)
+            laterRequests = argumentRequests[1:]
+            for argL, newContext, arg in self.enumeration(context, environment, argRequest,
+                                                          upperBound=upperBound,
+                                                          lowerBound=0.,
+                                                          maximumDepth=maximumDepth):
+                if violatesSymmetry(originalFunction, arg, argumentIndex):
+                    continue
+
+                newFunction = Application(function, arg)
+                for resultL, resultK, result in self.enumerateApplication(newContext, environment, newFunction,
+                                                                          laterRequests,
+                                                                          upperBound=upperBound + argL,
+                                                                          lowerBound=lowerBound + argL,
+                                                                          maximumDepth=maximumDepth,
+                                                                          originalFunction=originalFunction,
+                                                                          argumentIndex=argumentIndex + 1):
+                    yield resultL + argL, resultK, result
+
+    def sketchEnumeration(self,context,environment,request,sk,upperBound,
+                           maximumDepth=20,
+                           lowerBound=0.):
+        '''Enumerates all sketch instantiations whose MDL satisfies: lowerBound <= MDL < upperBound'''
+        if upperBound < 0. or maximumDepth == 1:
+            return
+
+        if sk.isHole:
+            yield from self.enumeration(context, environment, request, upperBound,
+                                        maximumDepth=maximumDepth,
+                                        lowerBound=lowerBound)
+        elif request.isArrow():
+            assert sk.isAbstraction
+            v = request.arguments[0]
+            for l, newContext, b in self.sketchEnumeration(context, [v] + environment,
+                                                           request.arguments[1],
+                                                           sk.body,
+                                                           upperBound=upperBound,
+                                                           lowerBound=lowerBound,
+                                                           maximumDepth=maximumDepth):
+                yield l, newContext, Abstraction(b)
+
+        else:
+            f, xs = sk.applicationParse()
+            if f.isIndex:
+                ft = environment[f.i].apply(context)
+            elif f.isInvented or f.isPrimitive:
+                context, ft = f.tp.instantiate(context)
+            elif f.isAbstraction:
+                assert False, "sketch is not in beta longform"
+            elif f.isHole:
+                assert False, "hole as function not yet supported"
+            elif f.isApplication:
+                assert False, "should never happen - bug in applicationParse"
+            else: assert False
+
+            try: context = context.unify(ft.returns(), request)                
+            except UnificationFailure:
+                print("Exception: sketch is ill-typed")
+                return #so that we can continue evaluating
+                # raise SketchEnumerationFailure() #"sketch is ill-typed"
+            ft = ft.apply(context)
+            argumentRequests = ft.functionArguments()
+
+            assert len(argumentRequests) == len(xs)
+
+            yield from self.sketchApplication(context, environment,
+                                              f, xs, argumentRequests,
+                                              upperBound=upperBound,
+                                              lowerBound=lowerBound,
+                                              maximumDepth=maximumDepth - 1)
+
+
+    def sketchApplication(self, context, environment,
+                          function, arguments, argumentRequests,
+                          # Upper bound on the description length of all of
+                          # the arguments
+                          upperBound,
+                          # Lower bound on the description length of all of
+                          # the arguments
+                          lowerBound=0.,
+                          maximumDepth=20):
+        if upperBound < 0. or maximumDepth == 1:
+            return
+
+        if argumentRequests == []:
+            if lowerBound <= 0. and 0. < upperBound:
+                yield 0., context, function
+            else:
+                return
+        else:
+            argRequest = argumentRequests[0].apply(context)
+            laterRequests = argumentRequests[1:]
+            firstSketch = arguments[0]
+            laterSketches = arguments[1:]
+            for argL, newContext, arg in self.sketchEnumeration(context, environment, argRequest,
+                                                                firstSketch,
+                                                                upperBound=upperBound,
+                                                                lowerBound=0.,
+                                                                maximumDepth=maximumDepth):
+
+                newFunction = Application(function, arg)
+                for resultL, resultK, result in self.sketchApplication(newContext, environment, newFunction,
+                                                                       laterSketches, laterRequests,
+                                                                       upperBound=upperBound + argL,
+                                                                       lowerBound=lowerBound + argL,
+                                                                       maximumDepth=maximumDepth):
+
+                    yield resultL + argL, resultK, result
+
+    def sketchLogLikelihood(self, request, full, sk, context=Context.EMPTY, environment=[]):
+        """
+        calculates mdl of full program 'full' from sketch 'sk'
+        """
+        if sk.isHole:
+            _, summary = self.likelihoodSummary(context, environment, request, full)
+            if summary is None:
+                eprint(
+                    "FATAL: program [ %s ] does not have a likelihood summary." %
+                    full, "r = ", request, "\n", self)
+                assert False
+            return summary.logLikelihood(self), context
+
+        elif request.isArrow():
+            assert sk.isAbstraction and full.isAbstraction
+            #assert sk.f == full.f #is this right? or do i need to recurse?
+            v = request.arguments[0]
+            return self.sketchLogLikelihood(request.arguments[1], full.body, sk.body, context=context, environment=[v] + environment)
+
+        else:
+            sk_f, sk_xs = sk.applicationParse()
+            full_f, full_xs = full.applicationParse()
+            if sk_f.isIndex:
+                assert sk_f == full_f, "sketch and full program don't match on an index"
+                ft = environment[sk_f.i].apply(context)
+            elif sk_f.isInvented or sk_f.isPrimitive:
+                assert sk_f == full_f, "sketch and full program don't match on a primitive"
+                context, ft = sk_f.tp.instantiate(context)
+            elif sk_f.isAbstraction:
+                assert False, "sketch is not in beta longform"
+            elif sk_f.isHole:
+                assert False, "hole as function not yet supported"
+            elif sk_f.isApplication:
+                assert False, "should never happen - bug in applicationParse"
+            else: assert False
+
+            try: context = context.unify(ft.returns(), request)                
+            except UnificationFailure: assert False, "sketch is ill-typed"
+            ft = ft.apply(context)
+            argumentRequests = ft.functionArguments()
+
+            assert len(argumentRequests) == len(sk_xs) == len(full_xs)  #this might not be true if holes??
+
+            return self.sketchllApplication(context, environment,
+                                              sk_f, sk_xs, full_f, full_xs, argumentRequests)
+
+    def sketchllApplication(self, context, environment,
+                          sk_function, sk_arguments, full_function, full_arguments, argumentRequests):
+        if argumentRequests == []:
+                return torch.tensor([0.]).cuda(), context #does this make sense?
+        else:
+            argRequest = argumentRequests[0].apply(context)
+            laterRequests = argumentRequests[1:]
+
+            sk_firstSketch = sk_arguments[0]
+            full_firstSketch = full_arguments[0]
+            sk_laterSketches = sk_arguments[1:]
+            full_laterSketches = full_arguments[1:]
+
+            argL, newContext = self.sketchLogLikelihood(argRequest, full_firstSketch, sk_firstSketch, context=context, environment=environment)
+
+            #finish this...
+            sk_newFunction = Application(sk_function, sk_firstSketch)  # is this redundant? maybe 
+            full_newFunction = Application(full_function, full_firstSketch)
+
+            resultL, context = self.sketchllApplication(newContext, environment, sk_newFunction, sk_laterSketches,
+                                            full_newFunction, full_laterSketches, laterRequests)
+
+            return resultL + argL, context
+
+        
+    def enumerateNearby(self, request, expr, distance=3.0):
+        """Enumerate programs with local mutations in subtrees with small description length"""
+        if distance <= 0:
+            yield expr
+        else:
+            def mutations(tp, loss):
+                for l, _, expr in self.enumeration(
+                        Context.EMPTY, [], tp, distance - loss):
+                    yield expr, l
+            yield from Mutator(self, mutations).execute(expr, request)
+
+
+    def enumerateHoles(self, request, expr, k=3, return_obj=Hole):
+        """Enumerate programs with a single hole within mdl distance"""
+        #TODO: make it possible to enumerate sketches with multiple holes
+        def mutations(tp, loss, is_left_application=False):
+            """
+            to allow applications lhs to become a hole,  
+            remove the condition below and ignore all the is_left_application kwds 
+            """
+            if not is_left_application: 
+                yield return_obj(), 0
+        top_k = []
+        for expr, l in Mutator(self, mutations).execute(expr, request):
+            if len(top_k) > 0:
+                i, v = min(enumerate(top_k), key=lambda x:x[1][1])
+                if l > v[1]:
+                    if len(top_k) >= k:
+                        top_k[i] = (expr, l)
+                    else:
+                        top_k.append((expr, l))
+                elif len(top_k) < k:
+                    top_k.append((expr, l))
+            else:
+                top_k.append((expr, l))
+        return sorted(top_k, key=lambda x:-x[1])
+
+    def untorch(self):
+        return Grammar(self.logVariable.data.tolist()[0], 
+                       [ (l.data.tolist()[0], t, p)
+                         for l, t, p in self.productions],
+                       continuationType=self.continuationType)
+
+class LikelihoodSummary(object):
+    '''Summarizes the terms that will be used in a likelihood calculation'''
+
+    def __init__(self):
+        self.uses = {}
+        self.normalizers = {}
+        self.constant = 0.
+
+    def __str__(self):
+        return """LikelihoodSummary(constant = %f,
+uses = {%s},
+normalizers = {%s})""" % (self.constant,
+                          ", ".join(
+                              "%s: %d" % (k,
+                                          v) for k,
+                              v in self.uses.items()),
+                          ", ".join(
+                              "%s: %d" % (k,
+                                          v) for k,
+                              v in self.normalizers.items()))
+
+    def record(self, actual, possibles, constant=0.):
+        # Variables are all normalized to be $0
+        if isinstance(actual, Index):
+            actual = Index(0)
+
+        # Make it something that we can hash
+        possibles = frozenset(sorted(possibles, key=hash))
+
+        self.constant += constant
+        self.uses[actual] = self.uses.get(actual, 0) + 1
+        self.normalizers[possibles] = self.normalizers.get(possibles, 0) + 1
+
+    def join(self, other):
+        self.constant += other.constant
+        for k, v in other.uses.items():
+            self.uses[k] = self.uses.get(k, 0) + v
+        for k, v in other.normalizers.items():
+            self.normalizers[k] = self.normalizers.get(k, 0) + v
+
+    def logLikelihood(self, grammar):
+        return self.constant + \
+            sum(count * grammar.expression2likelihood[p] for p, count in self.uses.items()) - \
+            sum(count * lse([grammar.expression2likelihood[p] for p in ps])
+                for ps, count in self.normalizers.items())
+    def logLikelihood_overlyGeneral(self, grammar):
+        """Calculates log likelihood of this summary, given that the summary might refer to productions that don't occur in the grammar"""
+        return self.constant + \
+            sum(count * grammar.expression2likelihood[p] for p, count in self.uses.items()) - \
+            sum(count * lse([grammar.expression2likelihood.get(p,NEGATIVEINFINITY) for p in ps])
+                for ps, count in self.normalizers.items())        
+    def numerator(self, grammar):
+        return self.constant + \
+            sum(count * grammar.expression2likelihood[p] for p, count in self.uses.items())
+    def denominator(self, grammar):
+        return \
+            sum(count * lse([grammar.expression2likelihood[p] for p in ps])
+                for ps, count in self.normalizers.items())
+    def toUses(self):
+        from collections import Counter
+        
+        possibleVariables = sum( count if Index(0) in ps else 0
+                                 for ps, count in self.normalizers.items() )
+        actualVariables = self.uses.get(Index(0), 0.)
+        actualUses = {k: v
+                      for k, v in self.uses.items()
+                      if not k.isIndex }
+        possibleUses = dict(Counter(p
+                                    for ps, count in self.normalizers.items()
+                                    for p_ in ps
+                                    if not p_.isIndex
+                                    for p in [p_]*count ))
+        return Uses(possibleVariables, actualVariables,
+                    possibleUses, actualUses)
+
+
+class Uses(object):
+    '''Tracks uses of different grammar productions'''
+
+    def __init__(self, possibleVariables=0., actualVariables=0.,
+                 possibleUses={}, actualUses={}):
+        self.actualVariables = actualVariables
+        self.possibleVariables = possibleVariables
+        self.possibleUses = possibleUses
+        self.actualUses = actualUses
+
+    def __str__(self):
+        return "Uses(actualVariables = %f, possibleVariables = %f, actualUses = %s, possibleUses = %s)" %\
+            (self.actualVariables, self.possibleVariables, self.actualUses, self.possibleUses)
+
+    def __repr__(self): return str(self)
+
+    def __mul__(self, a):
+        return Uses(a * self.possibleVariables,
+                    a * self.actualVariables,
+                    {p: a * u for p, u in self.possibleUses.items()},
+                    {p: a * u for p, u in self.actualUses.items()})
+
+    def __imul__(self, a):
+        self.possibleVariables *= a
+        self.actualVariables *= a
+        for p in self.possibleUses:
+            self.possibleUses[p] *= a
+        for p in self.actualUses:
+            self.actualUses[p] *= a
+        return self
+
+    def __rmul__(self, a):
+        return self * a
+
+    def __radd__(self, o):
+        if o == 0:
+            return self
+        return self + o
+
+    def __add__(self, o):
+        if o == 0:
+            return self
+
+        def merge(x, y):
+            z = x.copy()
+            for k, v in y.items():
+                z[k] = v + x.get(k, 0.)
+            return z
+        return Uses(self.possibleVariables + o.possibleVariables,
+                    self.actualVariables + o.actualVariables,
+                    merge(self.possibleUses, o.possibleUses),
+                    merge(self.actualUses, o.actualUses))
+
+    def __iadd__(self, o):
+        self.possibleVariables += o.possibleVariables
+        self.actualVariables += o.actualVariables
+        for k, v in o.possibleUses.items():
+            self.possibleUses[k] = self.possibleUses.get(k, 0.) + v
+        for k, v in o.actualUses.items():
+            self.actualUses[k] = self.actualUses.get(k, 0.) + v
+        return self
+
+    @staticmethod
+    def join(z, *weightedUses):
+        """Consumes weightedUses"""
+        if not weightedUses:
+            Uses.empty
+        if len(weightedUses) == 1:
+            return weightedUses[0][1]
+        for w, u in weightedUses:
+            u *= exp(w - z)
+        total = Uses()
+        total.possibleVariables = sum(
+            u.possibleVariables for _, u in weightedUses)
+        total.actualVariables = sum(u.actualVariables for _, u in weightedUses)
+        total.possibleUses = defaultdict(float)
+        total.actualUses = defaultdict(float)
+        for _, u in weightedUses:
+            for k, v in u.possibleUses.items():
+                total.possibleUses[k] += v
+            for k, v in u.actualUses.items():
+                total.actualUses[k] += v
+        return total
+
+
+Uses.empty = Uses()
+
+class ContextualGrammar:
+    def __init__(self, noParent, variableParent, library):
+        self.noParent, self.variableParent, self.library = noParent, variableParent, library
+
+        self.productions = [(None,t,p) for _,t,p in self.noParent.productions ]
+        self.primitives = [p for _,_2,p in self.productions ]
+
+        self.continuationType = noParent.continuationType
+        assert variableParent.continuationType == self.continuationType
+
+        assert set(noParent.primitives) == set(variableParent.primitives)
+        assert set(variableParent.primitives) == set(library.keys())
+        for e,gs in library.items():
+            assert len(gs) == len(e.infer().functionArguments())
+            for g in gs:
+                assert set(g.primitives) == set(library.keys())
+                assert g.continuationType == self.continuationType
+
+    def untorch(self):
+        return ContextualGrammar(self.noParent.untorch(), self.variableParent.untorch(),
+                                 {e: [g.untorch() for g in gs ]
+                                  for e,gs in self.library.items() })
+
+    def randomWeights(self, r):
+        """returns a new grammar with random weights drawn from r. calls `r` w/ old weight"""
+        return ContextualGrammar(self.noParent.randomWeights(r),
+                                 self.variableParent.randomWeights(r),
+                                 {e: [g.randomWeights(r) for g in gs]
+                                  for e,gs in self.library.items() })
+    def __str__(self):
+        lines = ["No parent:",str(self.noParent),"",
+                 "Variable parent:",str(self.variableParent),"",
+                 ""]
+        for e,gs in self.library.items():
+            for j,g in enumerate(gs):
+                lines.extend(["Parent %s, argument index %s"%(e,j),
+                              str(g),
+                              ""])
+        return "\n".join(lines)
+
+    def json(self):
+        return {"noParent": self.noParent.json(),
+                "variableParent": self.variableParent.json(),
+                "productions": [{"program": str(e),
+                                 "arguments": [gp.json() for gp in gs ]}
+                                    for e,gs in self.library.items() ]}
+
+    @staticmethod
+    def fromGrammar(g):
+        return ContextualGrammar(g, g,
+                                 {e: [g]*len(e.infer().functionArguments())
+                                  for e in g.primitives })
+                
+
+    class LS: # likelihood summary
+        def __init__(self, owner):
+            self.noParent = LikelihoodSummary()
+            self.variableParent = LikelihoodSummary()
+            self.library = {e: [LikelihoodSummary() for _ in gs]  for e,gs in owner.library.items() }
+
+        def record(self, parent, parentIndex, actual, possibles, constant):
+            if parent is None: ls = self.noParent
+            elif parent.isIndex: ls = self.variableParent
+            else: ls = self.library[parent][parentIndex]
+            ls.record(actual, possibles, constant=constant)
+
+        def join(self, other):
+            self.noParent.join(other.noParent)
+            self.variableParent.join(other.variableParent)
+            for e,gs in self.library.items():
+                for g1,g2 in zip(gs, other.library[e]):
+                    g1.join(g2)
+
+        def logLikelihood(self, owner):
+            return self.noParent.logLikelihood(owner.noParent) + \
+                   self.variableParent.logLikelihood(owner.variableParent) + \
+                   sum(r.logLikelihood(g)
+                       for e, rs in self.library.items()
+                       for r,g in zip(rs, owner.library[e]) )            
+        def numerator(self, owner):
+            return self.noParent.numerator(owner.noParent) + \
+                   self.variableParent.numerator(owner.variableParent) + \
+                   sum(r.numerator(g)
+                       for e, rs in self.library.items()
+                       for r,g in zip(rs, owner.library[e]) )            
+        def denominator(self, owner):
+            return self.noParent.denominator(owner.noParent) + \
+                   self.variableParent.denominator(owner.variableParent) + \
+                   sum(r.denominator(g)
+                       for e, rs in self.library.items()
+                       for r,g in zip(rs, owner.library[e]) )            
+
+    def likelihoodSummary(self, parent, parentIndex, context, environment, request, expression):
+        if request.isArrow():
+            assert expression.isAbstraction
+            return self.likelihoodSummary(parent, parentIndex,
+                                          context,
+                                          [request.arguments[0]] + environment,
+                                          request.arguments[1],
+                                          expression.body)
+        if parent is None: g = self.noParent
+        elif parent.isIndex: g = self.variableParent
+        else: g = self.library[parent][parentIndex]            
+        candidates = g.buildCandidates(request, context, environment,
+                                       normalize=False, returnTable=True)
+
+        # A list of everything that would have been possible to use here
+        possibles = [p for p in candidates.keys() if not p.isIndex]
+        numberOfVariables = sum(p.isIndex for p in candidates.keys())
+        if numberOfVariables > 0:
+            possibles += [Index(0)]
+
+        f, xs = expression.applicationParse()
+
+        assert f in candidates
+
+        thisSummary = self.LS(self)
+        thisSummary.record(parent, parentIndex,
+                           f, possibles,
+                           constant= -math.log(numberOfVariables) if f.isIndex else 0)
+
+        _, tp, context = candidates[f]
+        argumentTypes = tp.functionArguments()
+        assert len(xs) == len(argumentTypes)
+
+        for i, (argumentType, argument) in enumerate(zip(argumentTypes, xs)):
+            argumentType = argumentType.apply(context)
+            context, newSummary = self.likelihoodSummary(f, i,
+                                                         context, environment, argumentType, argument)
+            thisSummary.join(newSummary)
+
+        return context, thisSummary
+
+    def closedLikelihoodSummary(self, request, expression):
+        return self.likelihoodSummary(None,None,
+                                      Context.EMPTY,[],
+                                      request, expression)[1]
+
+    def logLikelihood(self, request, expression):
+        return self.closedLikelihoodSummary(request, expression).logLikelihood(self)
+
+    def sample(self, request, maximumDepth=8, maxAttempts=None):
+        attempts = 0
+        while True:
+            try:
+                _, e = self._sample(None, None, Context.EMPTY, [], request, maximumDepth)
+                return e
+            except NoCandidates:
+                if maxAttempts is not None:
+                    attempts += 1
+                    if attempts > maxAttempts: return None
+                continue
+            
+    def _sample(self, parent, parentIndex, context, environment, request, maximumDepth):
+        if request.isArrow():
+            context, body = self._sample(parent, parentIndex, context,
+                                         [request.arguments[0]] + environment,
+                                         request.arguments[1],
+                                         maximumDepth)
+            return context, Abstraction(body)
+        if parent is None: g = self.noParent
+        elif parent.isIndex: g = self.variableParent
+        else: g = self.library[parent][parentIndex]
+        candidates = g.buildCandidates(request, context, environment,
+                                       normalize=True, returnProbabilities=True,
+                                       mustBeLeaf=(maximumDepth <= 1))
+        newType, chosenPrimitive, context = sampleDistribution(candidates)
+
+        xs = newType.functionArguments()
+        returnValue = chosenPrimitive
+
+        for j,x in enumerate(xs):
+            x = x.apply(context)
+            context, x = self._sample(chosenPrimitive, j, context, environment, x, maximumDepth - 1)
+            returnValue = Application(returnValue, x)
+            
+        return context, returnValue
+
+    def expectedUsesMonteCarlo(self, request, debug=None):
+        import numpy as np
+        n = 0
+        u = [0.]*len(self.primitives)
+        primitives = list(sorted(self.primitives, key=str))
+        noInventions = all( not p.isInvented for p in primitives )
+        primitive2index = {primitive: i
+                           for i, primitive in enumerate(primitives)
+                           if primitive.isInvented or noInventions }
+        eprint(primitive2index)
+        ns = 10000
+        with timing(f"calculated expected uses using Monte Carlo simulation w/ {ns} samples"):
+            for _ in range(ns):
+                p = self.sample(request, maxAttempts=0)
+                if p is None: continue
+                n += 1
+                if debug and n < 10:
+                    eprint(debug, p)
+                for _, child in p.walk():
+                    if child not in primitive2index: continue
+                    u[primitive2index[child]] += 1.0
+        u = np.array(u)/n
+        if debug:
+            eprint(f"Got {n} samples. Feature vector:\n{u}")
+            eprint(f"Likely used primitives: {[p for p,i in primitive2index.items() if u[i] > 0.5]}")
+            eprint(f"Likely used primitive indices: {[i for p,i in primitive2index.items() if u[i] > 0.5]}")
+        return u
+
+    def featureVector(self, _=None, requests=None, onlyInventions=True, normalize=True):
+        """
+        Returns the probabilities licensed by the type system.
+        This is like the grammar productions, but with irrelevant junk removed.
+        Its intended use case is for clustering; it should be strictly better than the raw transition matrix.
+        """
+        if requests is None:
+            if self.continuationType: requests = {self.continuationType}
+            elif any( 'REAL' == str(p) for p in self.primitives ): requests = set()
+            elif any( 'STRING' == str(p) for p in self.primitives ): requests = {tlist(tcharacter)}
+            else: requests = set()
+        requests = {r.returns() for r in requests}
+        features = []
+        logWeights = []
+        for l,t,p in sorted(self.noParent.productions,
+                            key=lambda z: str(z[2])):
+            if onlyInventions and not p.isInvented: continue
+            if any( canUnify(r, t.returns()) for r in requests ) or len(requests) == 0:
+                logWeights.append(l)
+        features.append(logWeights)
+        for parent in sorted(self.primitives, key=str):
+            if onlyInventions and not parent.isInvented: continue
+            if parent not in self.library: continue
+            argumentTypes = parent.infer().functionArguments()
+            for j,g in enumerate(self.library[parent]):
+                argumentType = argumentTypes[j]
+                logWeights = []
+                for l,t,p in sorted(g.productions,
+                                    key=lambda z: str(z[2])):
+                    if onlyInventions and not p.isInvented: continue
+                    if canUnify(argumentType.returns(), t.returns()):
+                        logWeights.append(l)
+                features.append(logWeights)
+
+        if normalize:
+            features = [ [math.exp(w - z) for w in lw ]
+                         for lw in features
+                         if lw
+                         for z in [lse(lw)] ]
+        import numpy as np
+        return np.array([f
+                         for lw in features
+                         for f in lw])
+
+    def enumeration(self,context,environment,request,upperBound,
+                    parent=None, parentIndex=None,
+                    maximumDepth=20,
+                    lowerBound=0.):
+        '''Enumerates all programs whose MDL satisfies: lowerBound <= MDL < upperBound'''
+        if upperBound < 0 or maximumDepth == 1:
+            return
+
+        if request.isArrow():
+            v = request.arguments[0]
+            for l, newContext, b in self.enumeration(context, [v] + environment,
+                                                     request.arguments[1],
+                                                     parent=parent, parentIndex=parentIndex,
+                                                     upperBound=upperBound,
+                                                     lowerBound=lowerBound,
+                                                     maximumDepth=maximumDepth):
+                yield l, newContext, Abstraction(b)
+        else:
+            if parent is None: g = self.noParent
+            elif parent.isIndex: g = self.variableParent
+            else: g = self.library[parent][parentIndex]
+
+            candidates = g.buildCandidates(request, context, environment,
+                                           normalize=True)
+
+            for l, t, p, newContext in candidates:
+                mdl = -l
+                if not (mdl < upperBound):
+                    continue
+
+                xs = t.functionArguments()
+                for aL, aK, application in\
+                    self.enumerateApplication(newContext, environment, p, xs,
+                                              parent=p,
+                                              upperBound=upperBound + l,
+                                              lowerBound=lowerBound + l,
+                                              maximumDepth=maximumDepth - 1):
+                    yield aL + l, aK, application
+
+    def enumerateApplication(self, context, environment,
+                             function, argumentRequests,
+                             # Upper bound on the description length of all of
+                             # the arguments
+                             upperBound,
+                             # Lower bound on the description length of all of
+                             # the arguments
+                             lowerBound=0.,
+                             maximumDepth=20,
+                             parent=None, 
+                             originalFunction=None,
+                             argumentIndex=0):
+        assert parent is not None
+        if upperBound < 0. or maximumDepth == 1:
+            return
+        if originalFunction is None:
+            originalFunction = function
+
+        if argumentRequests == []:
+            if lowerBound <= 0. and 0. < upperBound:
+                yield 0., context, function
+            else:
+                return
+        else:
+            argRequest = argumentRequests[0].apply(context)
+            laterRequests = argumentRequests[1:]
+            for argL, newContext, arg in self.enumeration(context, environment, argRequest,
+                                                          parent=parent, parentIndex=argumentIndex,
+                                                          upperBound=upperBound,
+                                                          lowerBound=0.,
+                                                          maximumDepth=maximumDepth):
+                if violatesSymmetry(originalFunction, arg, argumentIndex):
+                    continue
+
+                newFunction = Application(function, arg)
+                for resultL, resultK, result in self.enumerateApplication(newContext, environment, newFunction,
+                                                                          laterRequests,
+                                                                          parent=parent,
+                                                                          upperBound=upperBound + argL,
+                                                                          lowerBound=lowerBound + argL,
+                                                                          maximumDepth=maximumDepth,
+                                                                          originalFunction=originalFunction,
+                                                                          argumentIndex=argumentIndex + 1):
+                    yield resultL + argL, resultK, result
+                
+        
+
+
+def violatesSymmetry(f, x, argumentIndex):
+    if not f.isPrimitive:
+        return False
+    while x.isApplication:
+        x = x.f
+    if not x.isPrimitive:
+        return False
+    f = f.name
+    x = x.name
+    if f == "car":
+        return x == "cons" or x == "empty"
+    if f == "cdr":
+        return x == "cons" or x == "empty"
+    if f == "+":
+        return x == "0" or (argumentIndex == 1 and x == "+")
+    if f == "-":
+        return argumentIndex == 1 and x == "0"
+    if f == "empty?":
+        return x == "cons" or x == "empty"
+    if f == "zero?":
+        return x == "0" or x == "1"
+    if f == "index" or f == "map" or f == "zip":
+        return x == "empty"
+    if f == "range":
+        return x == "0"
+    if f == "fold":
+        return argumentIndex == 1 and x == "empty"
+    return False
+
+def batchLikelihood(jobs):
+    """Takes as input a set of (program, request, grammar) and returns a dictionary mapping each of these to its likelihood under the grammar"""
+    superGrammar = Grammar.uniform(list({p for _1,_2,g in jobs for p in g.primitives}),
+                                   continuationType=list(jobs)[0][-1].continuationType)
+    programsAndRequests = {(program, request)
+                           for program, request, grammar in jobs}
+    with timing(f"Calculated {len(programsAndRequests)} likelihood summaries"):
+        summary = {(program, request): superGrammar.closedLikelihoodSummary(request, program)
+                   for program, request in programsAndRequests}
+    with timing(f"Calculated log likelihoods from summaries"):
+        response = {}
+        for program, request, grammar in jobs:
+            fast = summary[(program, request)].logLikelihood_overlyGeneral(grammar)
+            if False: # debugging
+                slow = grammar.logLikelihood(request, program)
+                print(program)
+                eprint(grammar.closedLikelihoodSummary(request, program))
+                eprint(superGrammar.closedLikelihoodSummary(request, program))
+                print()
+                assert abs(fast - slow) < 0.0001
+            response[(program, request, grammar)] = fast
+    return response
+
+if __name__ == "__main__":
+    from dreamcoder.domains.arithmetic.arithmeticPrimitives import *
+    g = ContextualGrammar.fromGrammar(Grammar.uniform([k0,k1,addition, subtraction]))
+    g = g.randomWeights(lambda *a: random.random())
+    #p = Program.parse("(lambda (+ 1 $0))")
+    request = arrow(tint,tint)
+    for ll,_,p in g.enumeration(Context.EMPTY,[],request,
+                               12.):
+        ll_ = g.logLikelihood(request,p)
+        print(ll,p,ll_)
+        d = abs(ll - ll_)
+        assert d < 0.0001

dreamcoder/likelihoodModel.py

@@ -0,0 +1,407 @@
+from dreamcoder.task import Task, EvaluationTimeout
+import gc
+from dreamcoder.utilities import *
+from collections import Counter
+import math
+
+from dreamcoder.domains.regex.groundtruthRegexes import gt_dict
+
+gt_dict = {"Data column no. "+str(num): r_str for num, r_str in gt_dict.items()}
+
+class AllOrNothingLikelihoodModel:
+    def __init__(self, timeout=None):
+        self.timeout = timeout
+
+    def score(self, program, task):
+        logLikelihood = task.logLikelihood(program, self.timeout)
+        return valid(logLikelihood), logLikelihood
+
+
+class EuclideanLikelihoodModel:
+    """Likelihood is based on Euclidean distance between features"""
+
+    def __init__(self, featureExtractor, successCutoff=0.9):
+        self.extract = featureExtractor
+        self.successCutoff = successCutoff
+
+    def score(self, program, task):
+        taskFeat = self.extract.featuresOfTask(task)
+        progFeat = self.extract.featuresOfProgram(program, task.request)
+        assert len(taskFeat) == len(progFeat)
+        distance = sum((x1 - x2)**2 for x1, x2 in zip(taskFeat, progFeat))
+        logLikelihood = float(-distance)  # FIXME: this is really naive
+        return exp(logLikelihood) > self.successCutoff, logLikelihood
+
+def longest_common_substr(arr):
+    #array of examples 
+
+# Python 3 program to find the stem
+# of given list of words
+# function to find the stem (longest
+# common substring) from the string array
+    # Determine size of the array
+    n = len(arr)
+
+    # Take first word from array
+    # as reference
+    s = arr[0]
+    l = len(s)
+    res = ""
+    for i in range(l) :
+        for j in range( i + 1, l + 1) :
+            # generating all possible substrings
+            # of our reference string arr[0] i.e s
+            stem = s[i:j]
+            k = 1
+            for k in range(1, n):
+
+                # Check if the generated stem is
+                # common to to all words
+                if stem not in arr[k]:
+                    break
+
+                # If current substring is present in
+                # all strings and its length is greater
+                # than current result
+            if (k + 1 == n and len(res) < len(stem)): res = stem
+    return res 
+
+def add_string_constants(tasks):
+    for task in tasks:
+        task.str_const = longest_common_substr([example[1] for example in task.examples])
+    return tasks
+
+def get_gt_ll(name, examples):
+    #gets groundtruth from dict
+    import pregex as pre
+    r_str = gt_dict[name]
+    preg = pre.create(r_str)
+
+    if type(examples[0]) == list:
+        examples = [ "".join(example) for example in examples]
+
+    s = sum( preg.match(example) for example in examples)
+    if s == float("-inf"):
+        print("bad for ", name)
+        print('preg:', preg)
+        print('preg sample:', [preg.sample() for i in range(3)])
+        print("exs", examples)
+        #assert False 
+    return s
+
+
+def add_cutoff_values(tasks, ll_cutoff):
+    from dreamcoder.domains.regex.makeRegexTasks import makeNewTasks
+    if ll_cutoff is None or ll_cutoff == "None":
+        for task in tasks:
+            task.ll_cutoff = None
+        return tasks
+    if ll_cutoff == "gt":
+        from dreamcoder.domains.regex.makeRegexTasks import regexHeldOutExamples
+        for task in tasks:
+            task.ll_cutoff = None
+            task.gt = get_gt_ll(task.name, [example[1] for example in task.examples])
+            task.gt_test = get_gt_ll(task.name,
+                                     [example[1] for example in regexHeldOutExamples(task) ])
+        return tasks
+    elif ll_cutoff == "plus":
+        for task in tasks:
+            task.ll_cutoff = regex_plus_bound([example[1] for example in task.examples])
+        return tasks
+    elif ll_cutoff == "bigram":
+        eprint("WARNING: using entire corpus to make bigram model")
+        #this means i do it twice, which is eh whatever
+        model = make_corpus_bigram(show_tasks(makeNewTasks()))
+        for task in tasks:
+            task.ll_cutoff = bigram_corpus_score([example[1] for example in task.examples], model)
+        return tasks
+    elif ll_cutoff =="unigram":
+        eprint("WARNING: using entire corpus to make unigram model")
+        #this means i do it twice, which is eh whatever
+        model = make_corpus_unigram(show_tasks(makeNewTasks()))
+        for task in tasks:
+            task.ll_cutoff = unigram_corpus_score([example[1] for example in task.examples], model)
+        return tasks
+    elif ll_cutoff =="mix":
+        eprint("WARNING: using entire corpus to make bigram model")
+        eprint("WARNING: using entire corpus to make unigram model")
+        #this means i do it twice, which is eh whatever
+        unigram = make_corpus_unigram(show_tasks(makeNewTasks()))
+        bigram = make_corpus_bigram(show_tasks(makeNewTasks()))
+        for task in tasks:
+            uniscore = unigram_corpus_score([example[1] for example in task.examples], unigram)
+            biscore = bigram_corpus_score([example[1] for example in task.examples], bigram)
+            task.ll_cutoff = math.log(0.75*math.exp(biscore) + 0.25*math.exp(uniscore))
+        return tasks
+    else:
+        eprint("not implemented")
+        eprint("cutoff val:")
+        eprint(ll_cutoff)
+        assert False
+
+def show_tasks(dataset):
+    task_list = []
+    for task in dataset:
+        task_list.append([example[1] for example in task.examples])
+    return task_list
+
+def regex_plus_bound(X):
+    from pregex import pregex
+    c = Counter(X)
+    regexes = [
+        pregex.create(".+"),
+        pregex.create("\d+"),
+        pregex.create("\w+"),
+        pregex.create("\s+"),
+        pregex.create("\\u+"),
+        pregex.create("\l+")]
+    regex_scores = []
+    for r in regexes:
+        regex_scores.append(sum(c[x] * r.match(x) for x in c)/float(sum([len(x) for x in X])) )
+    return max(regex_scores)
+
+
+def make_corpus_unigram(C):
+    str_list = [example + '\n' for task in C for example in task]
+    c = Counter(char for example in str_list for char in example )
+    n = sum(c.values())
+
+    logp = {x:math.log(c[x]/n) for x in c}
+    return logp
+
+def unigram_corpus_score(X, logp):
+    task_ll = 0
+    for x in X:
+        x = x + '\n'
+        task_ll += sum( logp.get(c, float('-inf')) for c in x)/len(x)
+
+    ll = task_ll/len(X)
+    return ll
+
+def unigram_task_score(X):
+    """
+    Given a list of strings, X, calculate the maximum log-likelihood per character for a unigram model over characters (including STOP symbol)
+    """
+    c = Counter(x for s in X for x in s)
+    c.update("end" for s in X)
+    n = sum(c.values())
+    logp = {x:math.log(c[x]/n) for x in c}
+    return sum(c[x]*logp[x] for x in c)/n
+
+def make_corpus_bigram(C):
+    #using newline as "end"
+    #C is a list of tasks
+
+    #make one big list of strings
+    str_list = [example + '\n' for task in C for example in task]
+
+    #make list of 
+    head_count = Counter(element[0] for element in str_list)
+    head_n = sum(head_count.values())
+    head_logp = {x:math.log(head_count[x]/head_n) for x in head_count}
+
+    body_count = Counter(element[i:i+2] for element in str_list for i in range(len(element)-1))
+    body_bigram_n = sum(body_count.values())
+    #body_count/body_bigram_n gives the joint of a bigram
+    body_character_n = Counter(char for element in str_list for char in element)
+    body_unigram_n = sum(body_character_n.values())
+
+    body_logp = {x:math.log(body_count[x] / body_bigram_n / body_character_n[x[0]] * body_unigram_n) for x in body_count}
+
+    return {**head_logp, **body_logp}
+
+def bigram_corpus_score(X, logp):
+    #assume you have a logp dict
+    task_ll = 0
+    for x in X:
+        bigram_list = [x[0]] + [x[i:i+2] for i in range(len(x)-1)] + [x[-1] + '\n']
+        bigram_list = [ ''.join(b) if isinstance(b,list) else b
+                        for b in bigram_list ]
+
+        string_ll = sum(logp.get(bigram, float('-inf')) for bigram in bigram_list) #/(len(x) + 1)
+
+        task_ll += string_ll
+
+    ll = task_ll #/len(X)
+    return ll
+
+
+class ProbabilisticLikelihoodModel:
+
+    def __init__(self, timeout):
+        self.timeout = timeout
+        # i need timeout
+
+    def score(self, program, task):
+        # need a try, catch here for problems, and for timeouts
+        # can copy task.py for the timeout structure
+        try:
+            def timeoutCallBack(_1, _2): raise EvaluationTimeout()
+            signal.signal(signal.SIGVTALRM, timeoutCallBack)
+            signal.setitimer(signal.ITIMER_VIRTUAL, self.timeout)
+            try:
+                string_pregex = program.evaluate([])
+                # if 'left_paren' in program.show(False):
+                #eprint("string_pregex:", string_pregex)
+                #eprint("string_pregex:", string_pregex)
+                preg = string_pregex  # pregex.create(string_pregex)
+            except IndexError:
+                # free variable
+                return False, NEGATIVEINFINITY
+            except Exception as e:
+                eprint("Exception during evaluation:", e)
+                if "Attempt to evaluate fragment variable" in e:
+                    eprint("program (bc fragment error)", program)
+                return False, NEGATIVEINFINITY
+
+        #tries and catches
+
+        # include prior somehow
+        # right now, just summing up log likelihoods. IDK if this is correct.
+        # also not using prior at all.
+
+            cum_ll = 0
+
+            example_list = [example[1] for example in task.examples]
+            c_example_list = Counter(example_list)
+
+            for c_example in c_example_list:
+                #might want a try, except around the following line:
+
+                try:
+                    #eprint("about to match", program)
+                    #print("preg:", preg)
+                    ll = preg.match(c_example)
+                    #eprint("completed match", ll, program)
+                except ValueError as e:
+                    eprint("ValueError:", e)
+                    ll = float('-inf')
+                
+                #eprint("pregex:", string_pregex)
+                #eprint("example[1]", example[1])
+
+                if ll == float('-inf'):
+                    return False, NEGATIVEINFINITY
+                else:
+                    #ll_per_char = ll/float(len(example[1]))
+                    #cum_ll_per_char += ll_per_char
+
+                    cum_ll += c_example_list[c_example] * ll
+            
+            #normalized_cum_ll_per_char = cum_ll_per_char/float(len(task.examples))
+            #avg_char_num = sum([len(example[1]) for example in task.examples])/float(len(task.examples))
+            
+            #cutoff_ll = regex_plus_bound(example_list)   
+
+            normalized_cum_ll = cum_ll/ float(sum([len(example) for example in example_list]))
+
+
+
+            #TODO: change the way normalized_cum_ll is calculated 
+            #TODO: refactor to pass in bigram_model, and others
+            #TODO: refactor to do 95% certainty thing josh wants
+            success = normalized_cum_ll > task.ll_cutoff
+
+
+
+            #eprint("cutoff_ll:", cutoff_ll, ", norm_cum_ll:", normalized_cum_ll)	
+
+            return success, normalized_cum_ll
+
+        except EvaluationTimeout:
+            eprint("Timed out while evaluating", program)
+            return False, NEGATIVEINFINITY
+        finally:
+            signal.signal(signal.SIGVTALRM, lambda *_: None)
+            signal.setitimer(signal.ITIMER_VIRTUAL, 0)
+
+
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
+    from torch.autograd import Variable
+
+    class FeatureDiscriminatorLikelihoodModel(nn.Module):
+        def __init__(self, tasks, featureExtractor,
+                     successCutoff=0.6, H=8, trainingSuccessRatio=0.5):
+            super(FeatureDiscriminatorLikelihoodModel, self).__init__()
+            self.extract = featureExtractor
+            self.successCutoff = successCutoff
+            self.trainingSuccessRatio = trainingSuccessRatio
+
+            self.W = nn.Linear(featureExtractor.outputDimensionality, H)
+            self.output = nn.Linear(H, 1)
+
+            # training on initialization
+            self.train(tasks)
+
+        def forward(self, examples):
+            """
+            Examples is a list of feature sets corresponding to a particular example.
+            Output in [0,1] whether all examples correspond to the same program
+            """
+            assert all(
+                len(x) == self.extract.outputDimensionality for x in examples)
+            examples = [F.tanh(self.W(ex)) for ex in examples]
+            maxed, _ = torch.max(torch.stack(examples), dim=0)
+            return F.sigmoid(self.output(maxed))
+
+        def train(self, tasks, steps=400):
+            # list of list of features for each example in each task
+            optimizer = torch.optim.Adam(self.parameters())
+            with timing("Trained discriminator"):
+                losses = []
+                for i in range(steps):
+                    self.zero_grad()
+                    if random.random() <= self.trainingSuccessRatio:
+                        # success
+                        t = random.choice(tasks)
+                        features = [self.extract.featuresOfTask(
+                            Task(t.name, t.request, [ex], t.features))
+                            for ex in t.examples]
+                        loss = (self(features) - 1.0)**2
+                    else:
+                        # fail
+                        t1, t2 = random.sample(tasks, 2)
+                        features1 = [self.extract.featuresOfTask(
+                            Task(t1.name, t1.request, [ex], t1.features))
+                            for ex in t1.examples[:len(t1.examples) / 2]]
+                        features2 = [self.extract.featuresOfTask(
+                            Task(t2.name, t2.request, [ex], t2.features))
+                            for ex in t2.examples[len(t2.examples) / 2:]]
+                        features = features1 + features2
+                        loss = self(features)**2
+
+                    loss.backward()
+                    optimizer.step()
+                    losses.append(loss.data[0])
+                    if not i % 50:
+                        eprint(
+                            "Discriminator Epoch",
+                            i,
+                            "Loss",
+                            sum(losses) /
+                            len(losses))
+                        gc.collect()
+
+        def score(self, program, task):
+            taskFeatures = self.extract.featuresOfTask(task)
+            progFeatures = self.extract.featuresOfProgram(
+                program, task.request)
+            likelihood = self([taskFeatures] + [progFeatures])
+            likelihood = float(likelihood)
+            return likelihood > self.successCutoff, log(likelihood)
+except ImportError:
+    pass
+
+
+if __name__=="__main__":
+
+    arr = ['MAM.OSBS.2014.06', 'MAM.OSBS.2013.07', 'MAM.OSBS.2013.09', 'MAM.OSBS.2014.05', 'MAM.OSBS.2014.11']
+    stems = longest_common_substr(arr)
+    print(stems)
+
+
+

dreamcoder/primitiveGraph.py

@@ -0,0 +1,182 @@
+from dreamcoder.program import *
+
+def graphPrimitives(result, prefix, view=False):
+    try:
+        from graphviz import Digraph
+    except:
+        eprint("You are missing the graphviz library - cannot graph primitives!")
+        return
+    
+
+    primitives = { p
+                   for g in result.grammars
+                   for p in g.primitives
+                   if p.isInvented }
+    age = {p: min(j for j,g in enumerate(result.grammars) if p in g.primitives) + 1
+           for p in primitives }
+
+
+
+    ages = set(age.values())
+    age2primitives = {a: {p for p,ap in age.items() if a == ap }
+                      for a in ages}
+
+    def lb(s,T=20):
+        s = s.split()
+        l = []
+        n = 0
+        for w in s:
+            if n + len(w) > T:
+                l.append("<br />")
+                n = 0
+            n += len(w)
+            l.append(w)
+        return " ".join(l)
+
+    nameSimplification = {
+        "fix1": 'Y',
+        "tower_loopM": "for",
+        "tower_embed": "get/set",
+        "moveHand": "move",
+        "reverseHand": "reverse",
+        "logo_DIVA": '/',
+        "logo_epsA": 'ε',
+        "logo_epsL": 'ε',
+        "logo_IFTY": '∞',
+        "logo_forLoop": "for",
+        "logo_UA": "2π",
+        "logo_FWRT": "move",
+        "logo_UL": "1",
+        "logo_SUBA": "-",
+        "logo_ZL": "0",
+        "logo_ZA": "0",
+        "logo_MULL": "*",
+        "logo_MULA": "*",
+        "logo_PT": "pen-up",
+        "logo_GETSET": "get/set"
+    }
+
+                
+    name = {}
+    simplification = {}
+    depth = {}
+    def getName(p):
+        if p in name: return name[p]
+        children = {k: getName(k)
+                    for _,k in p.body.walk()
+                    if k.isInvented}
+        simplification_ = p.body
+        for k,childName in children.items():
+            simplification_ = simplification_.substitute(k, Primitive(childName,None,None))
+        for original, simplified in nameSimplification.items():
+            simplification_ = simplification_.substitute(Primitive(original,None,None),
+                                                         Primitive(simplified,None,None))
+        name[p] = "f%d"%len(name)
+        simplification[p] = name[p] + '=' + lb(prettyProgram(simplification_, Lisp=True))
+        depth[p] = 1 + max([depth[k] for k in children] + [0])
+        return name[p]
+
+    for p in primitives:
+        getName(p)
+
+    depths = {depth[p] for p in primitives}
+    depth2primitives = {d: {p for p in primitives if depth[p] == d }
+                        for d in depths}
+
+    englishDescriptions = {"#(lambda (lambda (map (lambda (index $0 $2)) (range $0))))":
+                           "Prefix",
+                           "#(lambda (lambda (fold $0 $1 (lambda (lambda (cons $1 $0))))))":
+                           "Append",
+                           "#(lambda (cons LPAREN (#(lambda (lambda (fold $0 $1 (lambda (lambda (cons $1 $0)))))) (cons RPAREN empty) $0)))":
+                           "Enclose w/ parens",
+                           "#(lambda (unfold $0 (lambda (empty? $0)) (lambda (car $0)) (lambda (#(lambda (lambda (fold $1 $1 (lambda (lambda (cdr (if (char-eq? $1 $2) $3 $0))))))) $0 SPACE))))":
+                           "Abbreviate",
+                           "#(lambda (lambda (fold $1 $1 (lambda (lambda (cdr (if (char-eq? $1 $2) $3 $0)))))))":
+                           "Drop until char",
+                           "#(lambda (lambda (fold $1 $1 (lambda (lambda (if (char-eq? $1 $2) empty (cons $1 $0)))))))":
+                           "Take until char",
+                           "#(lambda (lambda (#(lambda (lambda (fold $0 $1 (lambda (lambda (cons $1 $0)))))) (cons $0 $1))))":
+                           "Append char",
+                           "#(lambda (lambda (map (lambda (if (char-eq? $0 $1) $2 $0)))))":
+                           "Substitute char",
+                           "#(lambda (lambda (length (unfold $1 (lambda (char-eq? (car $0) $1)) (lambda ',') (lambda (cdr $0))))))":
+                           "Index of char",
+                           "#(lambda (#(lambda (lambda (fold $0 $1 (lambda (lambda (cons $1 $0)))))) $0 STRING))":
+                           "Append const",
+                           "#(lambda (lambda (fold $1 $1 (lambda (lambda (fold $0 $0 (lambda (lambda (cdr (if (char-eq? $1 $4) $0 (cons $1 $0)))))))))))":
+                           "Last word",
+                           "#(lambda (lambda (cons (car $1) (cons '.' (cons (car $0) (cons '.' empty))))))":
+                           "Abbreviate name",
+                           "#(lambda (lambda (cons (car $1) (cons $0 empty))))":
+                           "First char+char",
+                           "#(lambda (#(lambda (lambda (fold $0 $1 (lambda (lambda (cons $1 $0)))))) (#(lambda (lambda (fold $1 $1 (lambda (lambda (fold $0 $0 (lambda (lambda (cdr (if (char-eq? $1 $4) $0 (cons $1 $0))))))))))) STRING (index (length (cdr $0)) $0)) $0))":
+                           "Ensure suffix"
+                           
+    }
+
+    def makeUnorderedGraph(fn):
+        g = Digraph()
+        g.graph_attr['rankdir'] = 'LR'
+
+        for p in primitives:
+            g.node(getName(p),
+                   label="<%s>"%simplification[p])
+        for p in primitives:
+            children = {k
+                        for _,k in p.body.walk()
+                        if k.isInvented}
+            for k in children:
+                g.edge(name[k],name[p])
+        try:
+            g.render(fn,view=view)
+            eprint("Exported primitive graph to",fn)
+        except:
+            eprint("Got some kind of error while trying to render primitive graph! Did you install graphviz/dot?")
+
+        
+
+    def makeGraph(ordering, fn):
+        g = Digraph()
+        g.graph_attr['rankdir'] = 'RL'
+
+        if False:
+            with g.subgraph(name='cluster_0') as sg:
+                sg.graph_attr['rank'] = 'same'
+                sg.attr(label='Primitives')
+                for j, primitive in enumerate(result.grammars[-1].primitives):
+                    if primitive.isInvented: continue
+                    sg.node("primitive%d"%j, label=str(primitive))
+
+        for o in sorted(ordering.keys()):
+            with g.subgraph(name='cluster_%d'%o) as sg:
+                sg.graph_attr['rank'] = 'same'
+                #sg.attr(label='Depth %d'%o)
+                for p in ordering[o]:
+                    if str(p) in englishDescriptions:
+                        thisLabel = '<<font face="boldfontname"><u>%s</u></font><br />%s>'%(englishDescriptions[str(p)],simplification[p])
+                    else:
+                        eprint("WARNING: Do not have an English description of:\n",p)
+                        eprint()
+                        thisLabel = "<%s>"%simplification[p]
+                    sg.node(getName(p),
+                            label=thisLabel)
+
+            for p in ordering[o]:
+                children = {k
+                            for _,k in p.body.walk()
+                            if k.isInvented}
+                for k in children:
+                    g.edge(name[k],name[p])
+
+        eprint("Exporting primitive graph to",fn)
+        try:
+            g.render(fn,view=view)
+        except Exception as e:
+            eprint("Got some kind of error while trying to render primitive graph! Did you install graphviz/dot?")
+            print(e)
+        
+        
+
+    makeGraph(depth2primitives,prefix+'depth.pdf')
+    makeUnorderedGraph(prefix+'unordered.pdf')
+    #makeGraph(age2primitives,prefix+'iter.pdf')

dreamcoder/program.py

@@ -0,0 +1,1214 @@
+# -*- coding: utf-8 -*-
+
+from dreamcoder.type import *
+from dreamcoder.utilities import *
+
+from time import time
+import math
+
+
+class InferenceFailure(Exception):
+    pass
+
+
+class ShiftFailure(Exception):
+    pass
+
+class RunFailure(Exception):
+    pass
+
+
+class Program(object):
+    def __repr__(self): return str(self)
+
+    def __ne__(self, o): return not (self == o)
+
+    def __str__(self): return self.show(False)
+
+    def canHaveType(self, t):
+        try:
+            context, actualType = self.inferType(Context.EMPTY, [], {})
+            context, t = t.instantiate(context)
+            context.unify(t, actualType)
+            return True
+        except UnificationFailure as e:
+            return False
+
+    def betaNormalForm(self):
+        n = self
+        while True:
+            np = n.betaReduce()
+            if np is None: return n
+            n = np
+
+    def infer(self):
+        try:
+            return self.inferType(Context.EMPTY, [], {})[1].canonical()
+        except UnificationFailure as e:
+            raise InferenceFailure(self, e)
+
+    def uncurry(self):
+        t = self.infer()
+        a = len(t.functionArguments())
+        e = self
+        existingAbstractions = 0
+        while e.isAbstraction:
+            e = e.body
+            existingAbstractions += 1
+        newAbstractions = a - existingAbstractions
+        assert newAbstractions >= 0
+
+        # e is the body stripped of abstractions. we are going to pile
+        # some more lambdas at the front, so free variables in e
+        # (which were bound to the stripped abstractions) need to be
+        # shifted by the number of abstractions that we will be adding
+        e = e.shift(newAbstractions)
+
+        for n in reversed(range(newAbstractions)):
+            e = Application(e, Index(n))
+        for _ in range(a):
+            e = Abstraction(e)
+
+        assert self.infer() == e.infer(), \
+            "FATAL: uncurry has a bug. %s : %s, but uncurried to %s : %s" % (self, self.infer(),
+                                                                             e, e.infer())
+        return e
+
+    def wellTyped(self):
+        try:
+            self.infer()
+            return True
+        except InferenceFailure:
+            return False
+
+    def runWithArguments(self, xs):
+        f = self.evaluate([])
+        for x in xs:
+            f = f(x)
+        return f
+
+    def applicationParses(self): yield self, []
+
+    def applicationParse(self): return self, []
+
+    @property
+    def closed(self):
+        for surroundingAbstractions, child in self.walk():
+            if isinstance(child, FragmentVariable):
+                return False
+            if isinstance(child, Index) and child.free(
+                    surroundingAbstractions):
+                return False
+        return True
+
+    @property
+    def numberOfFreeVariables(expression):
+        n = 0
+        for surroundingAbstractions, child in expression.walk():
+            # Free variable
+            if isinstance(child, Index) and child.free(
+                    surroundingAbstractions):
+                n = max(n, child.i - surroundingAbstractions + 1)
+        return n
+
+    def freeVariables(self):
+        for surroundingAbstractions, child in self.walk():
+            if child.isIndex and child.i >= surroundingAbstractions:
+                yield child.i - surroundingAbstractions
+
+    @property
+    def isIndex(self): return False
+
+    @property
+    def isUnion(self): return False
+
+    @property
+    def isApplication(self): return False
+
+    @property
+    def isAbstraction(self): return False
+
+    @property
+    def isPrimitive(self): return False
+
+    @property
+    def isInvented(self): return False
+
+    @property
+    def isHole(self): return False
+
+    @staticmethod
+    def parse(s):
+        s = parseSExpression(s)
+        def p(e):
+            if isinstance(e,list):
+                if e[0] == '#':
+                    assert len(e) == 2
+                    return Invented(p(e[1]))
+                if e[0] == 'lambda':
+                    assert len(e) == 2
+                    return Abstraction(p(e[1]))                    
+                f = p(e[0])
+                for x in e[1:]:
+                    f = Application(f,p(x))
+                return f
+            assert isinstance(e,str)
+            if e[0] == '$': return Index(int(e[1:]))
+            if e in Primitive.GLOBALS: return Primitive.GLOBALS[e]
+            if e == '??' or e == '?': return FragmentVariable.single
+            if e == '<HOLE>': return Hole.single
+            raise ParseFailure((s,e))
+        return p(s)
+
+    @staticmethod
+    def _parse(s,n):
+        while n < len(s) and s[n].isspace():
+            n += 1
+        for p in [
+                Application,
+                Abstraction,
+                Index,
+                Invented,
+                FragmentVariable,
+                Hole,
+                Primitive]:
+            try:
+                return p._parse(s,n)
+            except ParseFailure:
+                continue
+        raise ParseFailure(s)
+
+    # parser helpers
+    @staticmethod
+    def parseConstant(s,n,*constants):
+        for constant in constants:
+            try:
+                for i,c in enumerate(constant):
+                    if i + n >= len(s) or s[i + n] != c: raise ParseFailure(s)
+                return n + len(constant)
+            except ParseFailure: continue
+        raise ParseFailure(s)
+
+    @staticmethod
+    def parseHumanReadable(s):
+        s = parseSExpression(s)
+        def p(s, environment):
+            if isinstance(s, list) and s[0] in ['lambda','\\']:
+                assert isinstance(s[1], list) and len(s) == 3
+                newEnvironment = list(reversed(s[1])) + environment
+                e = p(s[2], newEnvironment)
+                for _ in s[1]: e = Abstraction(e)
+                return e
+            if isinstance(s, list):
+                a = p(s[0], environment)
+                for x in s[1:]:
+                    a = Application(a, p(x, environment))
+                return a
+            for j,v in enumerate(environment):
+                if s == v: return Index(j)
+            if s in Primitive.GLOBALS: return Primitive.GLOBALS[s]
+            assert False, f"could not parse {s}"
+        return p(s, [])
+                
+                
+
+
+class Application(Program):
+    '''Function application'''
+
+    def __init__(self, f, x):
+        self.f = f
+        self.x = x
+        self.hashCode = None
+        self.isConditional = (not isinstance(f,int)) and \
+                             f.isApplication and \
+                             f.f.isApplication and \
+                             f.f.f.isPrimitive and \
+                             f.f.f.name == "if"
+        if self.isConditional:
+            self.falseBranch = x
+            self.trueBranch = f.x
+            self.branch = f.f.x
+        else:
+            self.falseBranch = None
+            self.trueBranch = None
+            self.branch = None
+
+    def betaReduce(self):
+        # See if either the function or the argument can be reduced
+        f = self.f.betaReduce()
+        if f is not None: return Application(f,self.x)
+        x = self.x.betaReduce()
+        if x is not None: return Application(self.f,x)
+
+        # Neither of them could be reduced. Is this not a redex?
+        if not self.f.isAbstraction: return None
+
+        # Perform substitution
+        b = self.f.body
+        v = self.x
+        return b.substitute(Index(0), v.shift(1)).shift(-1)
+
+    def isBetaLong(self):
+        return (not self.f.isAbstraction) and self.f.isBetaLong() and self.x.isBetaLong()
+
+    def freeVariables(self):
+        return self.f.freeVariables() | self.x.freeVariables()
+
+    def clone(self): return Application(self.f.clone(), self.x.clone())
+
+    def annotateTypes(self, context, environment):
+        self.f.annotateTypes(context, environment)
+        self.x.annotateTypes(context, environment)
+        r = context.makeVariable()
+        context.unify(arrow(self.x.annotatedType, r), self.f.annotatedType)
+        self.annotatedType = r.applyMutable(context)        
+
+
+    @property
+    def isApplication(self): return True
+
+    def __eq__(
+        self,
+        other): return isinstance(
+        other,
+        Application) and self.f == other.f and self.x == other.x
+
+    def __hash__(self):
+        if self.hashCode is None:
+            self.hashCode = hash((hash(self.f), hash(self.x)))
+        return self.hashCode
+
+    """Because Python3 randomizes the hash function, we need to never pickle the hash"""
+    def __getstate__(self):
+        return self.f, self.x, self.isConditional, self.falseBranch, self.trueBranch, self.branch
+    def __setstate__(self, state):
+        try:
+            self.f, self.x, self.isConditional, self.falseBranch, self.trueBranch, self.branch = state
+        except ValueError:
+            # backward compatibility
+            assert 'x' in state
+            assert 'f' in state
+            f = state['f']
+            x = state['x']
+            self.f = f
+            self.x = x
+            self.isConditional = (not isinstance(f,int)) and \
+                                 f.isApplication and \
+                                 f.f.isApplication and \
+                                 f.f.f.isPrimitive and \
+                                 f.f.f.name == "if"
+            if self.isConditional:
+                self.falseBranch = x
+                self.trueBranch = f.x
+                self.branch = f.f.x
+            else:
+                self.falseBranch = None
+                self.trueBranch = None
+                self.branch = None
+
+        self.hashCode = None
+
+    def visit(self,
+              visitor,
+              *arguments,
+              **keywords): return visitor.application(self,
+                                                      *arguments,
+                                                      **keywords)
+
+    def show(self, isFunction):
+        if isFunction:
+            return "%s %s" % (self.f.show(True), self.x.show(False))
+        else:
+            return "(%s %s)" % (self.f.show(True), self.x.show(False))
+
+    def evaluate(self, environment):
+        if self.isConditional:
+            if self.branch.evaluate(environment):
+                return self.trueBranch.evaluate(environment)
+            else:
+                return self.falseBranch.evaluate(environment)
+        else:
+            return self.f.evaluate(environment)(self.x.evaluate(environment))
+
+    def inferType(self, context, environment, freeVariables):
+        (context, ft) = self.f.inferType(context, environment, freeVariables)
+        (context, xt) = self.x.inferType(context, environment, freeVariables)
+        (context, returnType) = context.makeVariable()
+        context = context.unify(ft, arrow(xt, returnType))
+        return (context, returnType.apply(context))
+
+    def applicationParses(self):
+        yield self, []
+        for f, xs in self.f.applicationParses():
+            yield f, xs + [self.x]
+
+    def applicationParse(self):
+        f, xs = self.f.applicationParse()
+        return f, xs + [self.x]
+
+    def shift(self, offset, depth=0):
+        return Application(self.f.shift(offset, depth),
+                           self.x.shift(offset, depth))
+
+    def substitute(self, old, new):
+        if self == old:
+            return new
+        return Application(
+            self.f.substitute(
+                old, new), self.x.substitute(
+                old, new))
+
+    def walkUncurried(self, d=0):
+        yield d, self
+        f, xs = self.applicationParse()
+        yield from f.walkUncurried(d)
+        for x in xs:
+            yield from x.walkUncurried(d)
+
+    def walk(self, surroundingAbstractions=0):
+        yield surroundingAbstractions, self
+        yield from self.f.walk(surroundingAbstractions)
+        yield from self.x.walk(surroundingAbstractions)
+
+    def size(self): return self.f.size() + self.x.size()
+
+    @staticmethod
+    def _parse(s,n):
+        while n < len(s) and s[n].isspace(): n += 1
+        if n == len(s) or s[n] != '(': raise ParseFailure(s)
+        n += 1
+
+        xs = []
+        while True:
+            x, n = Program._parse(s, n)
+            xs.append(x)
+            while n < len(s) and s[n].isspace(): n += 1
+            if n == len(s):
+                raise ParseFailure(s)
+            if s[n] == ")":
+                n += 1
+                break
+        e = xs[0]
+        for x in xs[1:]:
+            e = Application(e, x)
+        return e, n
+
+
+class Index(Program):
+    '''
+    deBruijn index: https://en.wikipedia.org/wiki/De_Bruijn_index
+    These indices encode variables.
+    '''
+
+    def __init__(self, i):
+        self.i = i
+
+    def show(self, isFunction): return "$%d" % self.i
+
+    def __eq__(self, o): return isinstance(o, Index) and o.i == self.i
+
+    def __hash__(self): return self.i
+
+    def visit(self,
+              visitor,
+              *arguments,
+              **keywords): return visitor.index(self,
+                                                *arguments,
+                                                **keywords)
+
+    def evaluate(self, environment):
+        return environment[self.i]
+
+    def inferType(self, context, environment, freeVariables):
+        if self.bound(len(environment)):
+            return (context, environment[self.i].apply(context))
+        else:
+            i = self.i - len(environment)
+            if i in freeVariables:
+                return (context, freeVariables[i].apply(context))
+            context, variable = context.makeVariable()
+            freeVariables[i] = variable
+            return (context, variable)
+
+    def clone(self): return Index(self.i)
+
+    def annotateTypes(self, context, environment):
+        self.annotatedType = environment[self.i].applyMutable(context)
+
+    def shift(self, offset, depth=0):
+        # bound variable
+        if self.bound(depth):
+            return self
+        else:  # free variable
+            i = self.i + offset
+            if i < 0:
+                raise ShiftFailure()
+            return Index(i)
+
+    def betaReduce(self): return None
+
+    def isBetaLong(self): return True
+
+    def freeVariables(self): return {self.i}
+
+    def substitute(self, old, new):
+        if old == self:
+            return new
+        else:
+            return self
+
+    def walk(self, surroundingAbstractions=0): yield surroundingAbstractions, self
+
+    def walkUncurried(self, d=0): yield d, self
+
+    def size(self): return 1
+
+    def free(self, surroundingAbstractions):
+        '''Is this index a free variable, given that it has surroundingAbstractions lambda's around it?'''
+        return self.i >= surroundingAbstractions
+
+    def bound(self, surroundingAbstractions):
+        '''Is this index a bound variable, given that it has surroundingAbstractions lambda's around it?'''
+        return self.i < surroundingAbstractions
+
+    @property
+    def isIndex(self): return True
+
+    @staticmethod
+    def _parse(s,n):
+        while n < len(s) and s[n].isspace(): n += 1
+        if n == len(s) or s[n] != '$':
+            raise ParseFailure(s)
+        n += 1
+        j = ""
+        while n < len(s) and s[n].isdigit():
+            j += s[n]
+            n += 1
+        if j == "":
+            raise ParseFailure(s)
+        return Index(int(j)), n
+
+
+class Abstraction(Program):
+    '''Lambda abstraction. Creates a new function.'''
+
+    def __init__(self, body):
+        self.body = body
+        self.hashCode = None
+
+    @property
+    def isAbstraction(self): return True
+
+    def __eq__(self, o): return isinstance(
+        o, Abstraction) and o.body == self.body
+
+    def __hash__(self):
+        if self.hashCode is None:
+            self.hashCode = hash((hash(self.body),))
+        return self.hashCode
+
+        """Because Python3 randomizes the hash function, we need to never pickle the hash"""
+    def __getstate__(self):
+        return self.body
+    def __setstate__(self, state):
+        self.body = state
+        self.hashCode = None
+
+    def isBetaLong(self): return self.body.isBetaLong()
+
+    def freeVariables(self):
+        return {f - 1 for f in self.body.freeVariables() if f > 0}
+
+    def visit(self,
+              visitor,
+              *arguments,
+              **keywords): return visitor.abstraction(self,
+                                                      *arguments,
+                                                      **keywords)
+
+    def clone(self): return Abstraction(self.body.clone())
+
+    def annotateTypes(self, context, environment):
+        v = context.makeVariable()
+        self.body.annotateTypes(context, [v] + environment)
+        self.annotatedType = arrow(v.applyMutable(context), self.body.annotatedType)
+
+    def show(self, isFunction):
+        return "(lambda %s)" % (self.body.show(False))
+
+    def evaluate(self, environment):
+        return lambda x: self.body.evaluate([x] + environment)
+
+    def betaReduce(self):
+        b = self.body.betaReduce()
+        if b is None: return None
+        return Abstraction(b)
+
+    def inferType(self, context, environment, freeVariables):
+        (context, argumentType) = context.makeVariable()
+        (context, returnType) = self.body.inferType(
+            context, [argumentType] + environment, freeVariables)
+        return (context, arrow(argumentType, returnType).apply(context))
+
+    def shift(self, offset, depth=0):
+        return Abstraction(self.body.shift(offset, depth + 1))
+
+    def substitute(self, old, new):
+        if self == old:
+            return new
+        old = old.shift(1)
+        new = new.shift(1)
+        return Abstraction(self.body.substitute(old, new))
+
+    def walk(self, surroundingAbstractions=0):
+        yield surroundingAbstractions, self
+        yield from self.body.walk(surroundingAbstractions + 1)
+
+    def walkUncurried(self, d=0):
+        yield d, self
+        yield from self.body.walkUncurried(d + 1)
+
+    def size(self): return self.body.size()
+
+    @staticmethod
+    def _parse(s,n):
+        n = Program.parseConstant(s,n,
+                                  '(\\','(lambda','(\u03bb')
+            
+        while n < len(s) and s[n].isspace(): n += 1
+
+        b, n = Program._parse(s,n)
+        while n < len(s) and s[n].isspace(): n += 1
+        n = Program.parseConstant(s,n,')')
+        return Abstraction(b), n
+
+
+class Primitive(Program):
+    GLOBALS = {}
+
+    def __init__(self, name, ty, value):
+        self.tp = ty
+        self.name = name
+        self.value = value
+        if name not in Primitive.GLOBALS:
+            Primitive.GLOBALS[name] = self
+
+    @property
+    def isPrimitive(self): return True
+
+    def __eq__(self, o): return isinstance(
+        o, Primitive) and o.name == self.name
+
+    def __hash__(self): return hash(self.name)
+
+    def visit(self,
+              visitor,
+              *arguments,
+              **keywords): return visitor.primitive(self,
+                                                    *arguments,
+                                                    **keywords)
+
+    def show(self, isFunction): return self.name
+
+    def clone(self): return Primitive(self.name, self.tp, self.value)
+
+    def annotateTypes(self, context, environment):
+        self.annotatedType = self.tp.instantiateMutable(context)
+
+    def evaluate(self, environment): return self.value
+
+    def betaReduce(self): return None
+
+    def isBetaLong(self): return True
+
+    def freeVariables(self): return set()
+
+    def inferType(self, context, environment, freeVariables):
+        return self.tp.instantiate(context)
+
+    def shift(self, offset, depth=0): return self
+
+    def substitute(self, old, new):
+        if self == old:
+            return new
+        else:
+            return self
+
+    def walk(self, surroundingAbstractions=0): yield surroundingAbstractions, self
+
+    def walkUncurried(self, d=0): yield d, self
+
+    def size(self): return 1
+
+    @staticmethod
+    def _parse(s,n):
+        while n < len(s) and s[n].isspace(): n += 1
+        name = []
+        while n < len(s) and not s[n].isspace() and s[n] not in '()':
+            name.append(s[n])
+            n += 1
+        name = "".join(name)
+        if name in Primitive.GLOBALS:
+            return Primitive.GLOBALS[name], n
+        raise ParseFailure(s)
+
+    # TODO(@mtensor): needs to be fixed to handle both pickling lambda functions and unpickling in general.
+    # def __getstate__(self):
+    #     return self.name
+
+    # def __setstate__(self, state):
+    #     #for backwards compatibility:
+    #     if type(state) == dict:
+    #         self.__dict__ = state
+    #     else:
+    #         p = Primitive.GLOBALS[state]
+    #         self.__init__(p.name, p.tp, p.value)
+
+class Invented(Program):
+    '''New invented primitives'''
+
+    def __init__(self, body):
+        self.body = body
+        self.tp = self.body.infer()
+        self.hashCode = None
+
+    @property
+    def isInvented(self): return True
+
+    def show(self, isFunction): return "#%s" % (self.body.show(False))
+
+    def visit(self,
+              visitor,
+              *arguments,
+              **keywords): return visitor.invented(self,
+                                                   *arguments,
+                                                   **keywords)
+
+    def __eq__(self, o): return isinstance(o, Invented) and o.body == self.body
+
+    def __hash__(self):
+        if self.hashCode is None:
+            self.hashCode = hash((0, hash(self.body)))
+        return self.hashCode
+
+    """Because Python3 randomizes the hash function, we need to never pickle the hash"""
+    def __getstate__(self):
+        return self.body, self.tp
+    def __setstate__(self, state):
+        self.body, self.tp = state
+        self.hashCode = None
+
+    def clone(self): return Invented(self.body)
+
+    def annotateTypes(self, context, environment):
+        self.annotatedType = self.tp.instantiateMutable(context)
+
+    def evaluate(self, e): return self.body.evaluate([])
+
+    def betaReduce(self): return self.body
+
+    def isBetaLong(self): return True
+
+    def freeVariables(self): return set()
+
+    def inferType(self, context, environment, freeVariables):
+        return self.tp.instantiate(context)
+
+    def shift(self, offset, depth=0): return self
+
+    def substitute(self, old, new):
+        if self == old:
+            return new
+        else:
+            return self
+
+    def walk(self, surroundingAbstractions=0): yield surroundingAbstractions, self
+
+    def walkUncurried(self, d=0): yield d, self
+
+    def size(self): return 1
+
+    @staticmethod
+    def _parse(s,n):
+        while n < len(s) and s[n].isspace(): n += 1
+        if n < len(s) and s[n] == '#':
+            n += 1
+            b,n = Program._parse(s,n)
+            return Invented(b),n
+        
+        raise ParseFailure(s)
+        
+
+class FragmentVariable(Program):
+    def __init__(self): pass
+
+    def show(self, isFunction): return "??"
+
+    def __eq__(self, o): return isinstance(o, FragmentVariable)
+
+    def __hash__(self): return 42
+
+    def visit(self, visitor, *arguments, **keywords):
+        return visitor.fragmentVariable(self, *arguments, **keywords)
+
+    def evaluate(self, e):
+        raise Exception('Attempt to evaluate fragment variable')
+
+    def betaReduce(self):
+        raise Exception('Attempt to beta reduce fragment variable')
+
+    def inferType(self, context, environment, freeVariables):
+        return context.makeVariable()
+
+    def shift(self, offset, depth=0):
+        raise Exception('Attempt to shift fragment variable')
+
+    def substitute(self, old, new):
+        if self == old:
+            return new
+        else:
+            return self
+
+    def match(
+            self,
+            context,
+            expression,
+            holes,
+            variableBindings,
+            environment=[]):
+        surroundingAbstractions = len(environment)
+        try:
+            context, variable = context.makeVariable()
+            holes.append(
+                (variable, expression.shift(-surroundingAbstractions)))
+            return context, variable
+        except ShiftFailure:
+            raise MatchFailure()
+
+    def walk(self, surroundingAbstractions=0): yield surroundingAbstractions, self
+
+    def walkUncurried(self, d=0): yield d, self
+
+    def size(self): return 1
+
+    @staticmethod
+    def _parse(s,n):
+        while n < len(s) and s[n].isspace(): n += 1
+        n = Program.parseConstant(s,n,'??','?')
+        return FragmentVariable.single, n
+
+FragmentVariable.single = FragmentVariable()
+
+
+class Hole(Program):
+    def __init__(self): pass
+
+    def show(self, isFunction): return "<HOLE>"
+
+    @property
+    def isHole(self): return True
+
+    def __eq__(self, o): return isinstance(o, Hole)
+
+    def __hash__(self): return 42
+
+    def evaluate(self, e):
+        raise Exception('Attempt to evaluate hole')
+
+    def betaReduce(self):
+        raise Exception('Attempt to beta reduce hole')
+
+    def inferType(self, context, environment, freeVariables):
+        return context.makeVariable()
+
+    def shift(self, offset, depth=0):
+        raise Exception('Attempt to shift fragment variable')
+
+    def walk(self, surroundingAbstractions=0): yield surroundingAbstractions, self
+
+    def walkUncurried(self, d=0): yield d, self
+
+    def size(self): return 1
+
+    @staticmethod
+    def _parse(s,n):
+        while n < len(s) and s[n].isspace(): n += 1
+        n = Program.parseConstant(s,n,
+                                  '<HOLE>')
+        return Hole.single, n
+
+
+Hole.single = Hole()
+
+
+class ShareVisitor(object):
+    def __init__(self):
+        self.primitiveTable = {}
+        self.inventedTable = {}
+        self.indexTable = {}
+        self.applicationTable = {}
+        self.abstractionTable = {}
+
+    def invented(self, e):
+        body = e.body.visit(self)
+        i = id(body)
+        if i in self.inventedTable:
+            return self.inventedTable[i]
+        new = Invented(body)
+        self.inventedTable[i] = new
+        return new
+
+    def primitive(self, e):
+        if e.name in self.primitiveTable:
+            return self.primitiveTable[e.name]
+        self.primitiveTable[e.name] = e
+        return e
+
+    def index(self, e):
+        if e.i in self.indexTable:
+            return self.indexTable[e.i]
+        self.indexTable[e.i] = e
+        return e
+
+    def application(self, e):
+        f = e.f.visit(self)
+        x = e.x.visit(self)
+        fi = id(f)
+        xi = id(x)
+        i = (fi, xi)
+        if i in self.applicationTable:
+            return self.applicationTable[i]
+        new = Application(f, x)
+        self.applicationTable[i] = new
+        return new
+
+    def abstraction(self, e):
+        body = e.body.visit(self)
+        i = id(body)
+        if i in self.abstractionTable:
+            return self.abstractionTable[i]
+        new = Abstraction(body)
+        self.abstractionTable[i] = new
+        return new
+
+    def execute(self, e):
+        return e.visit(self)
+
+
+class Mutator:
+    """Perform local mutations to an expr, yielding the expr and the
+    description length distance from the original program"""
+
+    def __init__(self, grammar, fn):
+        """Fn yields (expression, loglikelihood) from a type and loss.
+        Therefore, loss+loglikelihood is the distance from the original program."""
+        self.fn = fn
+        self.grammar = grammar
+        self.history = []
+
+    def enclose(self, expr):
+        for h in self.history[::-1]:
+            expr = h(expr)
+        return expr
+
+    def invented(self, e, tp, env, is_lhs=False):
+        deleted_ll = self.logLikelihood(tp, e, env)
+        for expr, replaced_ll in self.fn(tp, deleted, is_left_application=is_lhs):
+            yield self.enclose(expr), deleted_ll + replaced_ll
+
+    def primitive(self, e, tp, env, is_lhs=False):
+        deleted_ll = self.logLikelihood(tp, e, env)
+        for expr, replaced_ll in self.fn(tp, deleted_ll, is_left_application=is_lhs):
+            yield self.enclose(expr), deleted_ll + replaced_ll
+
+    def index(self, e, tp, env, is_lhs=False):
+        #yield from ()
+        deleted_ll = self.logLikelihood(tp, e, env) #self.grammar.logVariable
+        for expr, replaced_ll in self.fn(tp, deleted_ll, is_left_application=is_lhs):
+            yield self.enclose(expr), deleted_ll + replaced_ll
+
+    def application(self, e, tp, env, is_lhs=False):
+        self.history.append(lambda expr: Application(expr, e.x))
+        f_tp = arrow(e.x.infer(), tp)
+        yield from e.f.visit(self, f_tp, env, is_lhs=True)
+        self.history[-1] = lambda expr: Application(e.f, expr)
+        x_tp = inferArg(tp, e.f.infer())
+        yield from e.x.visit(self, x_tp, env)
+        self.history.pop()
+        deleted_ll = self.logLikelihood(tp, e, env)
+        for expr, replaced_ll in self.fn(tp, deleted_ll, is_left_application=is_lhs):
+            yield self.enclose(expr), deleted_ll + replaced_ll
+
+    def abstraction(self, e, tp, env, is_lhs=False):
+        self.history.append(lambda expr: Abstraction(expr))
+        yield from e.body.visit(self, tp.arguments[1], [tp.arguments[0]]+env)
+        self.history.pop()
+        deleted_ll = self.logLikelihood(tp, e, env)
+        for expr, replaced_ll in self.fn(tp, deleted_ll, is_left_application=is_lhs):
+            yield self.enclose(expr), deleted_ll + replaced_ll
+
+    def execute(self, e, tp):
+        yield from e.visit(self, tp, [])
+
+    def logLikelihood(self, tp, e, env):
+        summary = None
+        try:
+            _, summary = self.grammar.likelihoodSummary(Context.EMPTY, env,
+                tp, e, silent=True)
+        except AssertionError as err:
+            #print(f"closedLikelihoodSummary failed on tp={tp}, e={e}, error={err}")
+            pass
+        if summary is not None:
+            return summary.logLikelihood(self.grammar)
+        else:
+            tmpE, depth = e, 0
+            while isinstance(tmpE, Abstraction):
+                depth += 1
+                tmpE = tmpE.body
+            to_introduce = len(tp.functionArguments()) - depth
+            if to_introduce == 0:
+                #print(f"HIT NEGATIVEINFINITY, tp={tp}, e={e}")
+                return NEGATIVEINFINITY
+            for i in reversed(range(to_introduce)):
+                e = Application(e, Index(i))
+            for _ in range(to_introduce):
+                e = Abstraction(e)
+            return self.logLikelihood(tp, e, env)
+
+
+class RegisterPrimitives(object):
+    def invented(self, e): e.body.visit(self)
+
+    def primitive(self, e):
+        if e.name not in Primitive.GLOBALS:
+            Primitive(e.name, e.tp, e.value)
+
+    def index(self, e): pass
+
+    def application(self, e):
+        e.f.visit(self)
+        e.x.visit(self)
+
+    def abstraction(self, e): e.body.visit(self)
+
+    @staticmethod
+    def register(e): e.visit(RegisterPrimitives())
+
+
+class PrettyVisitor(object):
+    def __init__(self, Lisp=False):
+        self.Lisp = Lisp
+        self.numberOfVariables = 0
+        self.freeVariables = {}
+
+        self.variableNames = ["x", "y", "z", "u", "v", "w"]
+        self.variableNames += [chr(ord('a') + j)
+                               for j in range(20)]
+        self.toplevel = True
+
+    def makeVariable(self):
+        v = self.variableNames[self.numberOfVariables]
+        self.numberOfVariables += 1
+        return v
+
+    def invented(self, e, environment, isFunction, isAbstraction):
+        s = e.body.visit(self, [], isFunction, isAbstraction)
+        return s
+
+    def primitive(self, e, environment, isVariable, isAbstraction): return e.name
+
+    def index(self, e, environment, isVariable, isAbstraction):
+        if e.i < len(environment):
+            return environment[e.i]
+        else:
+            i = e.i - len(environment)
+            if i in self.freeVariables:
+                return self.freeVariables[i]
+            else:
+                v = self.makeVariable()
+                self.freeVariables[i] = v
+                return v
+
+    def application(self, e, environment, isFunction, isAbstraction):
+        self.toplevel = False
+        s = "%s %s" % (e.f.visit(self, environment, True, False),
+                       e.x.visit(self, environment, False, False))
+        if isFunction:
+            return s
+        else:
+            return "(" + s + ")"
+
+    def abstraction(self, e, environment, isFunction, isAbstraction):
+        toplevel = self.toplevel
+        self.toplevel = False
+        if not self.Lisp:
+            # Invent a new variable
+            v = self.makeVariable()
+            body = e.body.visit(self,
+                                [v] + environment,
+                                False,
+                                True)
+            if not e.body.isAbstraction:
+                body = "." + body
+            body = v + body
+            if not isAbstraction:
+                body = "λ" + body
+            if not toplevel:
+                body = "(%s)" % body
+            return body
+        else:
+            child = e
+            newVariables = []
+            while child.isAbstraction:
+                newVariables = [self.makeVariable()] + newVariables
+                child = child.body
+            body = child.visit(self, newVariables + environment,
+                               False, True)
+            body = "(λ (%s) %s)"%(" ".join(reversed(newVariables)), body)
+            return body
+            
+            
+
+def prettyProgram(e, Lisp=False):
+    return e.visit(PrettyVisitor(Lisp=Lisp), [], False, False)
+
+class EtaExpandFailure(Exception): pass
+class EtaLongVisitor(object):
+    """Converts an expression into eta-longform"""
+    def __init__(self, request=None):
+        self.request = request
+        self.context = None
+
+    def makeLong(self, e, request):
+        if request.isArrow():
+            # eta expansion
+            return Abstraction(Application(e.shift(1),
+                                           Index(0)))
+        return None
+        
+
+    def abstraction(self, e, request, environment):
+        if not request.isArrow(): raise EtaExpandFailure()
+        
+        return Abstraction(e.body.visit(self,
+                                        request.arguments[1],
+                                        [request.arguments[0]] + environment))
+
+    def _application(self, e, request, environment):
+        l = self.makeLong(e, request)
+        if l is not None: return l.visit(self, request, environment)
+
+        f, xs = e.applicationParse()
+
+        if f.isIndex:
+            ft = environment[f.i].applyMutable(self.context)
+        elif f.isInvented or f.isPrimitive:
+            ft = f.tp.instantiateMutable(self.context)
+        else: assert False, "Not in beta long form: %s"%e
+
+        self.context.unify(request, ft.returns())
+        ft = ft.applyMutable(self.context)
+
+        xt = ft.functionArguments()
+        if len(xs) != len(xt): raise EtaExpandFailure()
+
+        returnValue = f
+        for x,t in zip(xs,xt):
+            t = t.applyMutable(self.context)
+            returnValue = Application(returnValue,
+                                      x.visit(self, t, environment))
+        return returnValue
+
+    # This procedure works by recapitulating the generative process
+    # applications indices and primitives are all generated identically
+    
+    def application(self, e, request, environment): return self._application(e, request, environment)
+    
+    def index(self, e, request, environment): return self._application(e, request, environment)
+
+    def primitive(self, e, request, environment): return self._application(e, request, environment)
+
+    def invented(self, e, request, environment): return self._application(e, request, environment)
+
+    def execute(self, e):
+        assert len(e.freeVariables()) == 0
+        
+        if self.request is None:
+            eprint("WARNING: request not specified for etaexpansion")
+            self.request = e.infer()
+        self.context = MutableContext()
+        el = e.visit(self, self.request, [])
+        self.context = None
+        # assert el.infer().canonical() == e.infer().canonical(), \
+        #     f"Types are not preserved by ETA expansion: {e} : {e.infer().canonical()} vs {el} : {el.infer().canonical()}"
+        return el
+        
+
+
+class StripPrimitiveVisitor():
+    """Replaces all primitives .value's w/ None. Does not destructively modify anything"""
+    def invented(self,e):
+        return Invented(e.body.visit(self))
+    def primitive(self,e):
+        return Primitive(e.name,e.tp,None)
+    def application(self,e):
+        return Application(e.f.visit(self),
+                           e.x.visit(self))
+    def abstraction(self,e):
+        return Abstraction(e.body.visit(self))
+    def index(self,e): return e
+
+class ReplacePrimitiveValueVisitor():
+    """Intended to be used after StripPrimitiveVisitor.
+    Replaces all primitive.value's with their corresponding entry in Primitive.GLOBALS"""
+    def invented(self,e):
+        return Invented(e.body.visit(self))
+    def primitive(self,e):
+        return Primitive(e.name,e.tp,Primitive.GLOBALS[e.name].value)
+    def application(self,e):
+        return Application(e.f.visit(self),
+                           e.x.visit(self))
+    def abstraction(self,e):
+        return Abstraction(e.body.visit(self))
+    def index(self,e): return e
+
+def strip_primitive_values(e):
+    return e.visit(StripPrimitiveVisitor())
+def unstrip_primitive_values(e):
+    return e.visit(ReplacePrimitiveValueVisitor())
+    
+
+# from luke
+class TokeniseVisitor(object):
+    def invented(self, e):
+        return [e.body]
+
+    def primitive(self, e): return [e.name]
+
+    def index(self, e):
+        return ["$" + str(e.i)]
+
+    def application(self, e):
+        return ["("] + e.f.visit(self) + e.x.visit(self) + [")"]
+
+    def abstraction(self, e):
+        return ["(_lambda"] + e.body.visit(self) + [")_lambda"]
+
+
+def tokeniseProgram(e):
+    return e.visit(TokeniseVisitor())
+
+
+def untokeniseProgram(l):
+    lookup = {
+        "(_lambda": "(lambda",
+        ")_lambda": ")"
+    }
+    s = " ".join(lookup.get(x, x) for x in l)
+    return Program.parse(s)
+
+if __name__ == "__main__":
+    from dreamcoder.domains.arithmetic.arithmeticPrimitives import *
+    e = Program.parse("(#(lambda (?? (+ 1 $0))) (lambda (?? (+ 1 $0))) (lambda (?? (+ 1 $0))) - * (+ +))")
+    eprint(e)

dreamcoder/recognition.py

@@ -0,0 +1,1528 @@
+from dreamcoder.enumeration import *
+from dreamcoder.grammar import *
+# luke
+
+
+import gc
+
+try:
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    from torch.autograd import Variable
+    from torch.nn.utils.rnn import pack_padded_sequence
+except:
+    eprint("WARNING: Could not import torch. This is only okay when doing pypy compression.")
+
+try:
+    import numpy as np
+except:
+    eprint("WARNING: Could not import np. This is only okay when doing pypy compression.")
+    
+import json
+
+
+def variable(x, volatile=False, cuda=False):
+    if isinstance(x, list):
+        x = np.array(x)
+    if isinstance(x, (np.ndarray, np.generic)):
+        x = torch.from_numpy(x)
+    if cuda:
+        x = x.cuda()
+    return Variable(x, volatile=volatile)
+
+def maybe_cuda(x, use_cuda):
+    if use_cuda:
+        return x.cuda()
+    else:
+        return x
+
+
+def is_torch_not_a_number(v):
+    """checks whether a tortured variable is nan"""
+    v = v.data
+    if not ((v == v).item()):
+        return True
+    return False
+
+def is_torch_invalid(v):
+    """checks whether a torch variable is nan or inf"""
+    if is_torch_not_a_number(v):
+        return True
+    a = v - v
+    if is_torch_not_a_number(a):
+        return True
+    return False
+
+def _relu(x): return x.clamp(min=0)
+
+class Entropy(nn.Module):
+    """Calculates the entropy of logits"""
+    def __init__(self):
+        super(Entropy, self).__init__()
+
+    def forward(self, x):
+        b = F.softmax(x, dim=0) * F.log_softmax(x, dim=0)
+        b = -1.0 * b.sum()
+        return b
+
+class GrammarNetwork(nn.Module):
+    """Neural network that outputs a grammar"""
+    def __init__(self, inputDimensionality, grammar):
+        super(GrammarNetwork, self).__init__()
+        self.logProductions = nn.Linear(inputDimensionality, len(grammar)+1)
+        self.grammar = grammar
+        
+    def forward(self, x):
+        """Takes as input inputDimensionality-dimensional vector and returns Grammar
+        Tensor-valued probabilities"""
+        logProductions = self.logProductions(x)
+        return Grammar(logProductions[-1].view(1), #logVariable
+                       [(logProductions[k].view(1), t, program)
+                        for k, (_, t, program) in enumerate(self.grammar.productions)],
+                       continuationType=self.grammar.continuationType)
+
+    def batchedLogLikelihoods(self, xs, summaries):
+        """Takes as input BxinputDimensionality vector & B likelihood summaries;
+        returns B-dimensional vector containing log likelihood of each summary"""
+        use_cuda = xs.device.type == 'cuda'
+
+        B = xs.size(0)
+        assert len(summaries) == B
+        logProductions = self.logProductions(xs)
+
+        # uses[b][p] is # uses of primitive p by summary b
+        uses = np.zeros((B,len(self.grammar) + 1))
+        for b,summary in enumerate(summaries):
+            for p, production in enumerate(self.grammar.primitives):
+                uses[b,p] = summary.uses.get(production, 0.)
+            uses[b,len(self.grammar)] = summary.uses.get(Index(0), 0)
+
+        numerator = (logProductions * maybe_cuda(torch.from_numpy(uses).float(),use_cuda)).sum(1)
+        numerator += maybe_cuda(torch.tensor([summary.constant for summary in summaries ]).float(), use_cuda)
+
+        alternativeSet = {normalizer
+                          for s in summaries
+                          for normalizer in s.normalizers }
+        alternativeSet = list(alternativeSet)
+
+        mask = np.zeros((len(alternativeSet), len(self.grammar) + 1))
+        for tau in range(len(alternativeSet)):
+            for p, production in enumerate(self.grammar.primitives):
+                mask[tau,p] = 0. if production in alternativeSet[tau] else NEGATIVEINFINITY
+            mask[tau,len(self.grammar)] = 0. if Index(0) in alternativeSet[tau] else NEGATIVEINFINITY
+        mask = maybe_cuda(torch.tensor(mask).float(), use_cuda)
+
+        # mask: Rx|G|
+        # logProductions: Bx|G|
+        # Want: mask + logProductions : BxRx|G| = z
+        z = mask.repeat(B,1,1) + logProductions.repeat(len(alternativeSet),1,1).transpose(1,0)
+        # z: BxR
+        z = torch.logsumexp(z, 2) # pytorch 1.0 dependency
+
+        # Calculate how many times each normalizer was used
+        N = np.zeros((B, len(alternativeSet)))
+        for b, summary in enumerate(summaries):
+            for tau, alternatives in enumerate(alternativeSet):
+                N[b, tau] = summary.normalizers.get(alternatives,0.)
+
+        denominator = (maybe_cuda(torch.tensor(N).float(),use_cuda) * z).sum(1)
+        return numerator - denominator
+
+        
+
+class ContextualGrammarNetwork_LowRank(nn.Module):
+    def __init__(self, inputDimensionality, grammar, R=16):
+        """Low-rank approximation to bigram model. Parameters is linear in number of primitives.
+        R: maximum rank"""
+        
+        super(ContextualGrammarNetwork_LowRank, self).__init__()
+
+        self.grammar = grammar
+
+        self.R = R # embedding size
+
+        # library now just contains a list of indicies which go with each primitive
+        self.grammar = grammar
+        self.library = {}
+        self.n_grammars = 0
+        for prim in grammar.primitives:
+            numberOfArguments = len(prim.infer().functionArguments())
+            idx_list = list(range(self.n_grammars, self.n_grammars+numberOfArguments))
+            self.library[prim] = idx_list
+            self.n_grammars += numberOfArguments
+        
+        # We had an extra grammar for when there is no parent and for when the parent is a variable
+        self.n_grammars += 2
+        self.transitionMatrix = LowRank(inputDimensionality, self.n_grammars, len(grammar) + 1, R)
+        
+    def grammarFromVector(self, logProductions):
+        return Grammar(logProductions[-1].view(1),
+                       [(logProductions[k].view(1), t, program)
+                        for k, (_, t, program) in enumerate(self.grammar.productions)],
+                       continuationType=self.grammar.continuationType)
+
+    def forward(self, x):
+        assert len(x.size()) == 1, "contextual grammar doesn't currently support batching"
+
+        transitionMatrix = self.transitionMatrix(x)
+        
+        return ContextualGrammar(self.grammarFromVector(transitionMatrix[-1]), self.grammarFromVector(transitionMatrix[-2]),
+                {prim: [self.grammarFromVector(transitionMatrix[j]) for j in js]
+                 for prim, js in self.library.items()} )
+        
+    def vectorizedLogLikelihoods(self, x, summaries):
+        B = len(summaries)
+        G = len(self.grammar) + 1
+
+        # Which column of the transition matrix corresponds to which primitive
+        primitiveColumn = {p: c
+                           for c, (_1,_2,p) in enumerate(self.grammar.productions) }
+        primitiveColumn[Index(0)] = G - 1
+        # Which row of the transition matrix corresponds to which context
+        contextRow = {(parent, index): r
+                      for parent, indices in self.library.items()
+                      for index, r in enumerate(indices) }
+        contextRow[(None,None)] = self.n_grammars - 1
+        contextRow[(Index(0),None)] = self.n_grammars - 2
+
+        transitionMatrix = self.transitionMatrix(x)
+
+        # uses[b][g][p] is # uses of primitive p by summary b for parent g
+        uses = np.zeros((B,self.n_grammars,len(self.grammar)+1))
+        for b,summary in enumerate(summaries):
+            for e, ss in summary.library.items():
+                for g,s in zip(self.library[e], ss):
+                    assert g < self.n_grammars - 2
+                    for p, production in enumerate(self.grammar.primitives):
+                        uses[b,g,p] = s.uses.get(production, 0.)
+                    uses[b,g,len(self.grammar)] = s.uses.get(Index(0), 0)
+                    
+            # noParent: this is the last network output
+            for p, production in enumerate(self.grammar.primitives):            
+                uses[b, self.n_grammars - 1, p] = summary.noParent.uses.get(production, 0.)
+            uses[b, self.n_grammars - 1, G - 1] = summary.noParent.uses.get(Index(0), 0.)
+
+            # variableParent: this is the penultimate network output
+            for p, production in enumerate(self.grammar.primitives):            
+                uses[b, self.n_grammars - 2, p] = summary.variableParent.uses.get(production, 0.)
+            uses[b, self.n_grammars - 2, G - 1] = summary.variableParent.uses.get(Index(0), 0.)
+
+        uses = maybe_cuda(torch.tensor(uses).float(),use_cuda)
+        numerator = uses.view(B, -1) @ transitionMatrix.view(-1)
+        
+        constant = np.zeros(B)
+        for b,summary in enumerate(summaries):
+            constant[b] += summary.noParent.constant + summary.variableParent.constant
+            for ss in summary.library.values():
+                for s in ss:
+                    constant[b] += s.constant
+            
+        numerator = numerator + maybe_cuda(torch.tensor(constant).float(),use_cuda)
+
+        # Calculate the god-awful denominator
+        # Map from (parent, index, {set-of-alternatives}) to [occurrences-in-summary-zero, occurrences-in-summary-one, ...]
+        alternativeSet = {}
+        for b,summary in enumerate(summaries):
+            for normalizer, frequency in summary.noParent.normalizers.items():
+                k = (None,None,normalizer)
+                alternativeSet[k] = alternativeSet.get(k, np.zeros(B))
+                alternativeSet[k][b] += frequency
+            for normalizer, frequency in summary.variableParent.normalizers.items():
+                k = (Index(0),None,normalizer)
+                alternativeSet[k] = alternativeSet.get(k, np.zeros(B))
+                alternativeSet[k][b] += frequency
+            for parent, ss in summary.library.items():
+                for argumentIndex, s in enumerate(ss):
+                    for normalizer, frequency in s.normalizers.items():
+                        k = (parent, argumentIndex, normalizer)
+                        alternativeSet[k] = alternativeSet.get(k, zeros(B))
+                        alternativeSet[k][b] += frequency
+
+        # Calculate each distinct normalizing constant
+        alternativeNormalizer = {}
+        for parent, index, alternatives in alternativeSet:
+            r = transitionMatrix[contextRow[(parent, index)]]
+            entries = r[ [primitiveColumn[alternative] for alternative in alternatives ]]
+            alternativeNormalizer[(parent, index, alternatives)] = torch.logsumexp(entries, dim=0)
+
+        # Concatenate the normalizers into a vector
+        normalizerKeys = list(alternativeSet.keys())
+        normalizerVector = torch.cat([ alternativeNormalizer[k] for k in normalizerKeys])
+
+        assert False, "This function is still in progress."
+        
+
+    def batchedLogLikelihoods(self, xs, summaries):
+        """Takes as input BxinputDimensionality vector & B likelihood summaries;
+        returns B-dimensional vector containing log likelihood of each summary"""
+        use_cuda = xs.device.type == 'cuda'
+        
+        B = xs.shape[0]
+        G = len(self.grammar) + 1
+        assert len(summaries) == B
+
+        # logProductions: Bx n_grammars x G
+        logProductions = self.transitionMatrix(xs)
+        # uses[b][g][p] is # uses of primitive p by summary b for parent g
+        uses = np.zeros((B,self.n_grammars,len(self.grammar)+1))
+        for b,summary in enumerate(summaries):
+            for e, ss in summary.library.items():
+                for g,s in zip(self.library[e], ss):
+                    assert g < self.n_grammars - 2
+                    for p, production in enumerate(self.grammar.primitives):
+                        uses[b,g,p] = s.uses.get(production, 0.)
+                    uses[b,g,len(self.grammar)] = s.uses.get(Index(0), 0)
+                    
+            # noParent: this is the last network output
+            for p, production in enumerate(self.grammar.primitives):            
+                uses[b, self.n_grammars - 1, p] = summary.noParent.uses.get(production, 0.)
+            uses[b, self.n_grammars - 1, G - 1] = summary.noParent.uses.get(Index(0), 0.)
+
+            # variableParent: this is the penultimate network output
+            for p, production in enumerate(self.grammar.primitives):            
+                uses[b, self.n_grammars - 2, p] = summary.variableParent.uses.get(production, 0.)
+            uses[b, self.n_grammars - 2, G - 1] = summary.variableParent.uses.get(Index(0), 0.)
+            
+        numerator = (logProductions*maybe_cuda(torch.tensor(uses).float(),use_cuda)).view(B,-1).sum(1)
+
+        constant = np.zeros(B)
+        for b,summary in enumerate(summaries):
+            constant[b] += summary.noParent.constant + summary.variableParent.constant
+            for ss in summary.library.values():
+                for s in ss:
+                    constant[b] += s.constant
+            
+        numerator += maybe_cuda(torch.tensor(constant).float(),use_cuda)
+        
+        if True:
+
+            # Calculate the god-awful denominator
+            alternativeSet = set()
+            for summary in summaries:
+                for normalizer in summary.noParent.normalizers: alternativeSet.add(normalizer)
+                for normalizer in summary.variableParent.normalizers: alternativeSet.add(normalizer)
+                for ss in summary.library.values():
+                    for s in ss:
+                        for normalizer in s.normalizers: alternativeSet.add(normalizer)
+            alternativeSet = list(alternativeSet)
+
+            mask = np.zeros((len(alternativeSet), G))
+            for tau in range(len(alternativeSet)):
+                for p, production in enumerate(self.grammar.primitives):
+                    mask[tau,p] = 0. if production in alternativeSet[tau] else NEGATIVEINFINITY
+                mask[tau, G - 1] = 0. if Index(0) in alternativeSet[tau] else NEGATIVEINFINITY
+            mask = maybe_cuda(torch.tensor(mask).float(), use_cuda)
+
+            z = mask.repeat(self.n_grammars,1,1).repeat(B,1,1,1) + \
+                logProductions.repeat(len(alternativeSet),1,1,1).transpose(0,1).transpose(1,2)
+            z = torch.logsumexp(z, 3) # pytorch 1.0 dependency
+
+            N = np.zeros((B, self.n_grammars, len(alternativeSet)))
+            for b, summary in enumerate(summaries):
+                for e, ss in summary.library.items():
+                    for g,s in zip(self.library[e], ss):
+                        assert g < self.n_grammars - 2
+                        for r, alternatives in enumerate(alternativeSet):                
+                            N[b,g,r] = s.normalizers.get(alternatives, 0.)
+                # noParent: this is the last network output
+                for r, alternatives in enumerate(alternativeSet):
+                    N[b,self.n_grammars - 1,r] = summary.noParent.normalizers.get(alternatives, 0.)
+                # variableParent: this is the penultimate network output
+                for r, alternatives in enumerate(alternativeSet):
+                    N[b,self.n_grammars - 2,r] = summary.variableParent.normalizers.get(alternatives, 0.)
+            N = maybe_cuda(torch.tensor(N).float(), use_cuda)
+            denominator = (N*z).sum(1).sum(1)
+        else:
+            gs = [ self(xs[b]) for b in range(B) ]
+            denominator = torch.cat([ summary.denominator(g) for summary,g in zip(summaries, gs) ])
+            
+            
+
+        
+        
+        ll = numerator - denominator 
+
+        if False: # verifying that batching works correctly
+            gs = [ self(xs[b]) for b in range(B) ]
+            _l = torch.cat([ summary.logLikelihood(g) for summary,g in zip(summaries, gs) ])
+            assert torch.all((ll - _l).abs() < 0.0001)
+        return ll
+    
+class ContextualGrammarNetwork_Mask(nn.Module):
+    def __init__(self, inputDimensionality, grammar):
+        """Bigram model, but where the bigram transitions are unconditional.
+        Individual primitive probabilities are still conditional (predicted by neural network)
+        """
+        
+        super(ContextualGrammarNetwork_Mask, self).__init__()
+
+        self.grammar = grammar
+
+        # library now just contains a list of indicies which go with each primitive
+        self.grammar = grammar
+        self.library = {}
+        self.n_grammars = 0
+        for prim in grammar.primitives:
+            numberOfArguments = len(prim.infer().functionArguments())
+            idx_list = list(range(self.n_grammars, self.n_grammars+numberOfArguments))
+            self.library[prim] = idx_list
+            self.n_grammars += numberOfArguments
+        
+        # We had an extra grammar for when there is no parent and for when the parent is a variable
+        self.n_grammars += 2
+        self._transitionMatrix = nn.Parameter(nn.init.xavier_uniform(torch.Tensor(self.n_grammars, len(grammar) + 1)))
+        self._logProductions = nn.Linear(inputDimensionality, len(grammar)+1)
+
+    def transitionMatrix(self, x):
+        if len(x.shape) == 1: # not batched
+            return self._logProductions(x) + self._transitionMatrix # will broadcast
+        elif len(x.shape) == 2: # batched
+            return self._logProductions(x).unsqueeze(1).repeat(1,self.n_grammars,1) + \
+                self._transitionMatrix.unsqueeze(0).repeat(x.size(0),1,1)
+        else:
+            assert False, "unknown shape for transition matrix input"
+        
+    def grammarFromVector(self, logProductions):
+        return Grammar(logProductions[-1].view(1),
+                       [(logProductions[k].view(1), t, program)
+                        for k, (_, t, program) in enumerate(self.grammar.productions)],
+                       continuationType=self.grammar.continuationType)
+
+    def forward(self, x):
+        assert len(x.size()) == 1, "contextual grammar doesn't currently support batching"
+
+        transitionMatrix = self.transitionMatrix(x)
+        
+        return ContextualGrammar(self.grammarFromVector(transitionMatrix[-1]), self.grammarFromVector(transitionMatrix[-2]),
+                {prim: [self.grammarFromVector(transitionMatrix[j]) for j in js]
+                 for prim, js in self.library.items()} )
+        
+    def batchedLogLikelihoods(self, xs, summaries):
+        """Takes as input BxinputDimensionality vector & B likelihood summaries;
+        returns B-dimensional vector containing log likelihood of each summary"""
+        use_cuda = xs.device.type == 'cuda'
+        
+        B = xs.shape[0]
+        G = len(self.grammar) + 1
+        assert len(summaries) == B
+
+        # logProductions: Bx n_grammars x G
+        logProductions = self.transitionMatrix(xs)
+        # uses[b][g][p] is # uses of primitive p by summary b for parent g
+        uses = np.zeros((B,self.n_grammars,len(self.grammar)+1))
+        for b,summary in enumerate(summaries):
+            for e, ss in summary.library.items():
+                for g,s in zip(self.library[e], ss):
+                    assert g < self.n_grammars - 2
+                    for p, production in enumerate(self.grammar.primitives):
+                        uses[b,g,p] = s.uses.get(production, 0.)
+                    uses[b,g,len(self.grammar)] = s.uses.get(Index(0), 0)
+                    
+            # noParent: this is the last network output
+            for p, production in enumerate(self.grammar.primitives):            
+                uses[b, self.n_grammars - 1, p] = summary.noParent.uses.get(production, 0.)
+            uses[b, self.n_grammars - 1, G - 1] = summary.noParent.uses.get(Index(0), 0.)
+
+            # variableParent: this is the penultimate network output
+            for p, production in enumerate(self.grammar.primitives):            
+                uses[b, self.n_grammars - 2, p] = summary.variableParent.uses.get(production, 0.)
+            uses[b, self.n_grammars - 2, G - 1] = summary.variableParent.uses.get(Index(0), 0.)
+            
+        numerator = (logProductions*maybe_cuda(torch.tensor(uses).float(),use_cuda)).view(B,-1).sum(1)
+
+        constant = np.zeros(B)
+        for b,summary in enumerate(summaries):
+            constant[b] += summary.noParent.constant + summary.variableParent.constant
+            for ss in summary.library.values():
+                for s in ss:
+                    constant[b] += s.constant
+            
+        numerator += maybe_cuda(torch.tensor(constant).float(),use_cuda)
+        
+        if True:
+
+            # Calculate the god-awful denominator
+            alternativeSet = set()
+            for summary in summaries:
+                for normalizer in summary.noParent.normalizers: alternativeSet.add(normalizer)
+                for normalizer in summary.variableParent.normalizers: alternativeSet.add(normalizer)
+                for ss in summary.library.values():
+                    for s in ss:
+                        for normalizer in s.normalizers: alternativeSet.add(normalizer)
+            alternativeSet = list(alternativeSet)
+
+            mask = np.zeros((len(alternativeSet), G))
+            for tau in range(len(alternativeSet)):
+                for p, production in enumerate(self.grammar.primitives):
+                    mask[tau,p] = 0. if production in alternativeSet[tau] else NEGATIVEINFINITY
+                mask[tau, G - 1] = 0. if Index(0) in alternativeSet[tau] else NEGATIVEINFINITY
+            mask = maybe_cuda(torch.tensor(mask).float(), use_cuda)
+
+            z = mask.repeat(self.n_grammars,1,1).repeat(B,1,1,1) + \
+                logProductions.repeat(len(alternativeSet),1,1,1).transpose(0,1).transpose(1,2)
+            z = torch.logsumexp(z, 3) # pytorch 1.0 dependency
+
+            N = np.zeros((B, self.n_grammars, len(alternativeSet)))
+            for b, summary in enumerate(summaries):
+                for e, ss in summary.library.items():
+                    for g,s in zip(self.library[e], ss):
+                        assert g < self.n_grammars - 2
+                        for r, alternatives in enumerate(alternativeSet):                
+                            N[b,g,r] = s.normalizers.get(alternatives, 0.)
+                # noParent: this is the last network output
+                for r, alternatives in enumerate(alternativeSet):
+                    N[b,self.n_grammars - 1,r] = summary.noParent.normalizers.get(alternatives, 0.)
+                # variableParent: this is the penultimate network output
+                for r, alternatives in enumerate(alternativeSet):
+                    N[b,self.n_grammars - 2,r] = summary.variableParent.normalizers.get(alternatives, 0.)
+            N = maybe_cuda(torch.tensor(N).float(), use_cuda)
+            denominator = (N*z).sum(1).sum(1)
+        else:
+            gs = [ self(xs[b]) for b in range(B) ]
+            denominator = torch.cat([ summary.denominator(g) for summary,g in zip(summaries, gs) ])
+            
+            
+
+        
+        
+        ll = numerator - denominator
+
+        if False: # verifying that batching works correctly
+            gs = [ self(xs[b]) for b in range(B) ]
+            _l = torch.cat([ summary.logLikelihood(g) for summary,g in zip(summaries, gs) ])
+            assert torch.all((ll - _l).abs() < 0.0001)
+        return ll
+        
+                
+
+class ContextualGrammarNetwork(nn.Module):
+    """Like GrammarNetwork but ~contextual~"""
+    def __init__(self, inputDimensionality, grammar):
+        super(ContextualGrammarNetwork, self).__init__()
+        
+        # library now just contains a list of indicies which go with each primitive
+        self.grammar = grammar
+        self.library = {}
+        self.n_grammars = 0
+        for prim in grammar.primitives:
+            numberOfArguments = len(prim.infer().functionArguments())
+            idx_list = list(range(self.n_grammars, self.n_grammars+numberOfArguments))
+            self.library[prim] = idx_list
+            self.n_grammars += numberOfArguments
+        
+        # We had an extra grammar for when there is no parent and for when the parent is a variable
+        self.n_grammars += 2
+        self.network = nn.Linear(inputDimensionality, (self.n_grammars)*(len(grammar) + 1))
+
+
+    def grammarFromVector(self, logProductions):
+        return Grammar(logProductions[-1].view(1),
+                       [(logProductions[k].view(1), t, program)
+                        for k, (_, t, program) in enumerate(self.grammar.productions)],
+                       continuationType=self.grammar.continuationType)
+
+    def forward(self, x):
+        assert len(x.size()) == 1, "contextual grammar doesn't currently support batching"
+
+        allVars = self.network(x).view(self.n_grammars, -1)
+        return ContextualGrammar(self.grammarFromVector(allVars[-1]), self.grammarFromVector(allVars[-2]),
+                {prim: [self.grammarFromVector(allVars[j]) for j in js]
+                 for prim, js in self.library.items()} )
+
+    def batchedLogLikelihoods(self, xs, summaries):
+        use_cuda = xs.device.type == 'cuda'
+        """Takes as input BxinputDimensionality vector & B likelihood summaries;
+        returns B-dimensional vector containing log likelihood of each summary"""
+
+        B = xs.shape[0]
+        G = len(self.grammar) + 1
+        assert len(summaries) == B
+
+        # logProductions: Bx n_grammars x G
+        logProductions = self.network(xs).view(B, self.n_grammars, G)
+        # uses[b][g][p] is # uses of primitive p by summary b for parent g
+        uses = np.zeros((B,self.n_grammars,len(self.grammar)+1))
+        for b,summary in enumerate(summaries):
+            for e, ss in summary.library.items():
+                for g,s in zip(self.library[e], ss):
+                    assert g < self.n_grammars - 2
+                    for p, production in enumerate(self.grammar.primitives):
+                        uses[b,g,p] = s.uses.get(production, 0.)
+                    uses[b,g,len(self.grammar)] = s.uses.get(Index(0), 0)
+                    
+            # noParent: this is the last network output
+            for p, production in enumerate(self.grammar.primitives):            
+                uses[b, self.n_grammars - 1, p] = summary.noParent.uses.get(production, 0.)
+            uses[b, self.n_grammars - 1, G - 1] = summary.noParent.uses.get(Index(0), 0.)
+
+            # variableParent: this is the penultimate network output
+            for p, production in enumerate(self.grammar.primitives):            
+                uses[b, self.n_grammars - 2, p] = summary.variableParent.uses.get(production, 0.)
+            uses[b, self.n_grammars - 2, G - 1] = summary.variableParent.uses.get(Index(0), 0.)
+            
+        numerator = (logProductions*maybe_cuda(torch.tensor(uses).float(),use_cuda)).view(B,-1).sum(1)
+
+        constant = np.zeros(B)
+        for b,summary in enumerate(summaries):
+            constant[b] += summary.noParent.constant + summary.variableParent.constant
+            for ss in summary.library.values():
+                for s in ss:
+                    constant[b] += s.constant
+            
+        numerator += maybe_cuda(torch.tensor(constant).float(),use_cuda)
+
+        # Calculate the god-awful denominator
+        alternativeSet = set()
+        for summary in summaries:
+            for normalizer in summary.noParent.normalizers: alternativeSet.add(normalizer)
+            for normalizer in summary.variableParent.normalizers: alternativeSet.add(normalizer)
+            for ss in summary.library.values():
+                for s in ss:
+                    for normalizer in s.normalizers: alternativeSet.add(normalizer)
+        alternativeSet = list(alternativeSet)
+
+        mask = np.zeros((len(alternativeSet), G))
+        for tau in range(len(alternativeSet)):
+            for p, production in enumerate(self.grammar.primitives):
+                mask[tau,p] = 0. if production in alternativeSet[tau] else NEGATIVEINFINITY
+            mask[tau, G - 1] = 0. if Index(0) in alternativeSet[tau] else NEGATIVEINFINITY
+        mask = maybe_cuda(torch.tensor(mask).float(), use_cuda)
+
+        z = mask.repeat(self.n_grammars,1,1).repeat(B,1,1,1) + \
+            logProductions.repeat(len(alternativeSet),1,1,1).transpose(0,1).transpose(1,2)
+        z = torch.logsumexp(z, 3) # pytorch 1.0 dependency
+
+        N = np.zeros((B, self.n_grammars, len(alternativeSet)))
+        for b, summary in enumerate(summaries):
+            for e, ss in summary.library.items():
+                for g,s in zip(self.library[e], ss):
+                    assert g < self.n_grammars - 2
+                    for r, alternatives in enumerate(alternativeSet):                
+                        N[b,g,r] = s.normalizers.get(alternatives, 0.)
+            # noParent: this is the last network output
+            for r, alternatives in enumerate(alternativeSet):
+                N[b,self.n_grammars - 1,r] = summary.noParent.normalizers.get(alternatives, 0.)
+            # variableParent: this is the penultimate network output
+            for r, alternatives in enumerate(alternativeSet):
+                N[b,self.n_grammars - 2,r] = summary.variableParent.normalizers.get(alternatives, 0.)
+        N = maybe_cuda(torch.tensor(N).float(), use_cuda)
+        
+
+        
+        denominator = (N*z).sum(1).sum(1)
+        ll = numerator - denominator
+
+        if False: # verifying that batching works correctly
+            gs = [ self(xs[b]) for b in range(B) ]
+            _l = torch.cat([ summary.logLikelihood(g) for summary,g in zip(summaries, gs) ])
+            assert torch.all((ll - _l).abs() < 0.0001)
+
+        return ll
+        
+
+class RecognitionModel(nn.Module):
+    def __init__(self,featureExtractor,grammar,hidden=[64],activation="tanh",
+                 rank=None,contextual=False,mask=False,
+                 cuda=False,
+                 previousRecognitionModel=None,
+                 id=0):
+        super(RecognitionModel, self).__init__()
+        self.id = id
+        self.trained=False
+        self.use_cuda = cuda
+
+        self.featureExtractor = featureExtractor
+        # Sanity check - make sure that all of the parameters of the
+        # feature extractor were added to our parameters as well
+        if hasattr(featureExtractor, 'parameters'):
+            for parameter in featureExtractor.parameters():
+                assert any(myParameter is parameter for myParameter in self.parameters())
+
+        # Build the multilayer perceptron that is sandwiched between the feature extractor and the grammar
+        if activation == "sigmoid":
+            activation = nn.Sigmoid
+        elif activation == "relu":
+            activation = nn.ReLU
+        elif activation == "tanh":
+            activation = nn.Tanh
+        else:
+            raise Exception('Unknown activation function ' + str(activation))
+        self._MLP = nn.Sequential(*[ layer
+                                     for j in range(len(hidden))
+                                     for layer in [
+                                             nn.Linear(([featureExtractor.outputDimensionality] + hidden)[j],
+                                                       hidden[j]),
+                                             activation()]])
+
+        self.entropy = Entropy()
+
+        if len(hidden) > 0:
+            self.outputDimensionality = self._MLP[-2].out_features
+            assert self.outputDimensionality == hidden[-1]
+        else:
+            self.outputDimensionality = self.featureExtractor.outputDimensionality
+
+        self.contextual = contextual
+        if self.contextual:
+            if mask:
+                self.grammarBuilder = ContextualGrammarNetwork_Mask(self.outputDimensionality, grammar)
+            else:
+                self.grammarBuilder = ContextualGrammarNetwork_LowRank(self.outputDimensionality, grammar, rank)
+        else:
+            self.grammarBuilder = GrammarNetwork(self.outputDimensionality, grammar)
+        
+        self.grammar = ContextualGrammar.fromGrammar(grammar) if contextual else grammar
+        self.generativeModel = grammar
+        
+        self._auxiliaryPrediction = nn.Linear(self.featureExtractor.outputDimensionality, 
+                                              len(self.grammar.primitives))
+        self._auxiliaryLoss = nn.BCEWithLogitsLoss()
+
+        if cuda: self.cuda()
+
+        if previousRecognitionModel:
+            self._MLP.load_state_dict(previousRecognitionModel._MLP.state_dict())
+            self.featureExtractor.load_state_dict(previousRecognitionModel.featureExtractor.state_dict())
+            
+    def auxiliaryLoss(self, frontier, features):
+        # Compute a vector of uses
+        ls = frontier.bestPosterior.program
+        def uses(summary):
+            if hasattr(summary, 'uses'): 
+                return torch.tensor([ float(int(p in summary.uses))
+                                      for p in self.generativeModel.primitives ])
+            assert hasattr(summary, 'noParent')
+            u = uses(summary.noParent) + uses(summary.variableParent)
+            for ss in summary.library.values():
+                for s in ss:
+                    u += uses(s)
+            return u
+        u = uses(ls)
+        u[u > 1.] = 1.
+        if self.use_cuda: u = u.cuda()
+        al = self._auxiliaryLoss(self._auxiliaryPrediction(features), u)
+        return al
+            
+    def taskEmbeddings(self, tasks):
+        return {task: self.featureExtractor.featuresOfTask(task).data.cpu().numpy()
+                for task in tasks}
+
+    def forward(self, features):
+        """returns either a Grammar or a ContextualGrammar
+        Takes as input the output of featureExtractor.featuresOfTask"""
+        features = self._MLP(features)
+        return self.grammarBuilder(features)
+
+    def auxiliaryPrimitiveEmbeddings(self):
+        """Returns the actual outputDimensionality weight vectors for each of the primitives."""
+        auxiliaryWeights = self._auxiliaryPrediction.weight.data.cpu().numpy()
+        primitivesDict =  {self.grammar.primitives[i] : auxiliaryWeights[i, :] for i in range(len(self.grammar.primitives))}
+        return primitivesDict
+
+    def grammarOfTask(self, task):
+        features = self.featureExtractor.featuresOfTask(task)
+        if features is None: return None
+        return self(features)
+
+    def grammarLogProductionsOfTask(self, task):
+        """Returns the grammar logits from non-contextual models."""
+
+        features = self.featureExtractor.featuresOfTask(task)
+        if features is None: return None
+
+        if hasattr(self, 'hiddenLayers'):
+            # Backward compatability with old checkpoints.
+            for layer in self.hiddenLayers:
+                features = self.activation(layer(features))
+            # return features
+            return self.noParent[1](features)
+        else:
+            features = self._MLP(features)
+
+        if self.contextual:
+            if hasattr(self.grammarBuilder, 'variableParent'):
+                return self.grammarBuilder.variableParent.logProductions(features)
+            elif hasattr(self.grammarBuilder, 'network'):
+                return self.grammarBuilder.network(features).view(-1)
+            elif hasattr(self.grammarBuilder, 'transitionMatrix'):
+                return self.grammarBuilder.transitionMatrix(features).view(-1)
+            else:
+                assert False
+        else:
+            return self.grammarBuilder.logProductions(features)
+
+    def grammarFeatureLogProductionsOfTask(self, task):
+        return torch.tensor(self.grammarOfTask(task).untorch().featureVector())
+
+    def grammarLogProductionDistanceToTask(self, task, tasks):
+        """Returns the cosine similarity of all other tasks to a given task."""
+        taskLogits = self.grammarLogProductionsOfTask(task).unsqueeze(0) # Change to [1, D]
+        assert taskLogits is not None, 'Grammar log productions are not defined for this task.'
+        otherTasks = [t for t in tasks if t is not task] # [nTasks -1 , D]
+
+        # Build matrix of all other tasks.
+        otherLogits = torch.stack([self.grammarLogProductionsOfTask(t) for t in otherTasks])
+        cos = nn.CosineSimilarity(dim=1, eps=1e-6)
+        cosMatrix = cos(taskLogits, otherLogits)
+        return cosMatrix.data.cpu().numpy()
+
+    def grammarEntropyOfTask(self, task):
+        """Returns the entropy of the grammar distribution from non-contextual models for a task."""
+        grammarLogProductionsOfTask = self.grammarLogProductionsOfTask(task)
+
+        if grammarLogProductionsOfTask is None: return None
+
+        if hasattr(self, 'entropy'):
+            return self.entropy(grammarLogProductionsOfTask)
+        else:
+            e = Entropy()
+            return e(grammarLogProductionsOfTask)
+
+    def taskAuxiliaryLossLayer(self, tasks):
+        return {task: self._auxiliaryPrediction(self.featureExtractor.featuresOfTask(task)).view(-1).data.cpu().numpy()
+                for task in tasks}
+                
+    def taskGrammarFeatureLogProductions(self, tasks):
+        return {task: self.grammarFeatureLogProductionsOfTask(task).data.cpu().numpy()
+                for task in tasks}
+
+    def taskGrammarLogProductions(self, tasks):
+        return {task: self.grammarLogProductionsOfTask(task).data.cpu().numpy()
+                for task in tasks}
+
+    def taskGrammarStartProductions(self, tasks):
+        return {task: np.array([l for l,_1,_2 in g.productions ])
+                for task in tasks
+                for g in [self.grammarOfTask(task).untorch().noParent] }
+
+    def taskHiddenStates(self, tasks):
+        return {task: self._MLP(self.featureExtractor.featuresOfTask(task)).view(-1).data.cpu().numpy()
+                for task in tasks}
+
+    def taskGrammarEntropies(self, tasks):
+        return {task: self.grammarEntropyOfTask(task).data.cpu().numpy()
+                for task in tasks}
+
+    def frontierKL(self, frontier, auxiliary=False, vectorized=True):
+        features = self.featureExtractor.featuresOfTask(frontier.task)
+        if features is None:
+            return None, None
+        # Monte Carlo estimate: draw a sample from the frontier
+        entry = frontier.sample()
+
+        al = self.auxiliaryLoss(frontier, features if auxiliary else features.detach())
+
+        if not vectorized:
+            g = self(features)
+            return - entry.program.logLikelihood(g), al
+        else:
+            features = self._MLP(features).unsqueeze(0)
+            
+            ll = self.grammarBuilder.batchedLogLikelihoods(features, [entry.program]).view(-1)
+            return -ll, al
+            
+
+    def frontierBiasOptimal(self, frontier, auxiliary=False, vectorized=True):
+        if not vectorized:
+            features = self.featureExtractor.featuresOfTask(frontier.task)
+            if features is None: return None, None
+            al = self.auxiliaryLoss(frontier, features if auxiliary else features.detach())
+            g = self(features)
+            summaries = [entry.program for entry in frontier]
+            likelihoods = torch.cat([entry.program.logLikelihood(g) + entry.logLikelihood
+                                     for entry in frontier ])
+            best = likelihoods.max()
+            return -best, al
+            
+        batchSize = len(frontier.entries)
+        features = self.featureExtractor.featuresOfTask(frontier.task)
+        if features is None: return None, None
+        al = self.auxiliaryLoss(frontier, features if auxiliary else features.detach())
+        features = self._MLP(features)
+        features = features.expand(batchSize, features.size(-1))  # TODO
+        lls = self.grammarBuilder.batchedLogLikelihoods(features, [entry.program for entry in frontier])
+        actual_ll = torch.Tensor([ entry.logLikelihood for entry in frontier])
+        lls = lls + (actual_ll.cuda() if self.use_cuda else actual_ll)
+        ml = -lls.max() #Beware that inputs to max change output type
+        return ml, al
+
+    def replaceProgramsWithLikelihoodSummaries(self, frontier):
+        return Frontier(
+            [FrontierEntry(
+                program=self.grammar.closedLikelihoodSummary(frontier.task.request, e.program),
+                logLikelihood=e.logLikelihood,
+                logPrior=e.logPrior) for e in frontier],
+            task=frontier.task)
+
+    def train(self, frontiers, _=None, steps=None, lr=0.001, topK=5, CPUs=1,
+              timeout=None, evaluationTimeout=0.001,
+              helmholtzFrontiers=[], helmholtzRatio=0., helmholtzBatch=500,
+              biasOptimal=None, defaultRequest=None, auxLoss=False, vectorized=True):
+        """
+        helmholtzRatio: What fraction of the training data should be forward samples from the generative model?
+        helmholtzFrontiers: Frontiers from programs enumerated from generative model (optional)
+        If helmholtzFrontiers is not provided then we will sample programs during training
+        """
+        assert (steps is not None) or (timeout is not None), \
+            "Cannot train recognition model without either a bound on the number of gradient steps or bound on the training time"
+        if steps is None: steps = 9999999
+        if biasOptimal is None: biasOptimal = len(helmholtzFrontiers) > 0
+        
+        requests = [frontier.task.request for frontier in frontiers]
+        if len(requests) == 0 and helmholtzRatio > 0 and len(helmholtzFrontiers) == 0:
+            assert defaultRequest is not None, "You are trying to random Helmholtz training, but don't have any frontiers. Therefore we would not know the type of the program to sample. Try specifying defaultRequest=..."
+            requests = [defaultRequest]
+        frontiers = [frontier.topK(topK).normalize()
+                     for frontier in frontiers if not frontier.empty]
+        if len(frontiers) == 0:
+            eprint("You didn't give me any nonempty replay frontiers to learn from. Going to learn from 100% Helmholtz samples")
+            helmholtzRatio = 1.
+
+        # Should we sample programs or use the enumerated programs?
+        randomHelmholtz = len(helmholtzFrontiers) == 0
+        
+        class HelmholtzEntry:
+            def __init__(self, frontier, owner):
+                self.request = frontier.task.request
+                self.task = None
+                self.programs = [e.program for e in frontier]
+                self.frontier = Thunk(lambda: owner.replaceProgramsWithLikelihoodSummaries(frontier))
+                self.owner = owner
+
+            def clear(self): self.task = None
+
+            def calculateTask(self):
+                assert self.task is None
+                p = random.choice(self.programs)
+                return self.owner.featureExtractor.taskOfProgram(p, self.request)
+
+            def makeFrontier(self):
+                assert self.task is not None
+                f = Frontier(self.frontier.force().entries,
+                             task=self.task)
+                return f
+        
+            
+            
+
+        # Should we recompute tasks on the fly from Helmholtz?  This
+        # should be done if the task is stochastic, or if there are
+        # different kinds of inputs on which it could be run. For
+        # example, lists and strings need this; towers and graphics do
+        # not. There is no harm in recomputed the tasks, it just
+        # wastes time.
+        if not hasattr(self.featureExtractor, 'recomputeTasks'):
+            self.featureExtractor.recomputeTasks = True
+        helmholtzFrontiers = [HelmholtzEntry(f, self)
+                              for f in helmholtzFrontiers]
+        random.shuffle(helmholtzFrontiers)
+        
+        helmholtzIndex = [0]
+        def getHelmholtz():
+            if randomHelmholtz:
+                if helmholtzIndex[0] >= len(helmholtzFrontiers):
+                    updateHelmholtzTasks()
+                    helmholtzIndex[0] = 0
+                    return getHelmholtz()
+                helmholtzIndex[0] += 1
+                return helmholtzFrontiers[helmholtzIndex[0] - 1].makeFrontier()
+
+            f = helmholtzFrontiers[helmholtzIndex[0]]
+            if f.task is None:
+                with timing("Evaluated another batch of Helmholtz tasks"):
+                    updateHelmholtzTasks()
+                return getHelmholtz()
+
+            helmholtzIndex[0] += 1
+            if helmholtzIndex[0] >= len(helmholtzFrontiers):
+                helmholtzIndex[0] = 0
+                random.shuffle(helmholtzFrontiers)
+                if self.featureExtractor.recomputeTasks:
+                    for fp in helmholtzFrontiers:
+                        fp.clear()
+                    return getHelmholtz() # because we just cleared everything
+            assert f.task is not None
+            return f.makeFrontier()
+            
+        def updateHelmholtzTasks():
+            updateCPUs = CPUs if hasattr(self.featureExtractor, 'parallelTaskOfProgram') and self.featureExtractor.parallelTaskOfProgram else 1
+            if updateCPUs > 1: eprint("Updating Helmholtz tasks with",updateCPUs,"CPUs",
+                                      "while using",getThisMemoryUsage(),"memory")
+            
+            if randomHelmholtz:
+                newFrontiers = self.sampleManyHelmholtz(requests, helmholtzBatch, CPUs)
+                newEntries = []
+                for f in newFrontiers:
+                    e = HelmholtzEntry(f,self)
+                    e.task = f.task
+                    newEntries.append(e)
+                helmholtzFrontiers.clear()
+                helmholtzFrontiers.extend(newEntries)
+                return 
+
+            # Save some memory by freeing up the tasks as we go through them
+            if self.featureExtractor.recomputeTasks:
+                for hi in range(max(0, helmholtzIndex[0] - helmholtzBatch,
+                                    min(helmholtzIndex[0], len(helmholtzFrontiers)))):
+                    helmholtzFrontiers[hi].clear()
+
+            if hasattr(self.featureExtractor, 'tasksOfPrograms'):
+                eprint("batching task calculation")
+                newTasks = self.featureExtractor.tasksOfPrograms(
+                    [random.choice(hf.programs)
+                     for hf in helmholtzFrontiers[helmholtzIndex[0]:helmholtzIndex[0] + helmholtzBatch] ],
+                    [hf.request
+                     for hf in helmholtzFrontiers[helmholtzIndex[0]:helmholtzIndex[0] + helmholtzBatch] ])
+            else:
+                newTasks = [hf.calculateTask() 
+                            for hf in helmholtzFrontiers[helmholtzIndex[0]:helmholtzIndex[0] + helmholtzBatch]]
+
+                """
+                # catwong: Disabled for ensemble training.
+                newTasks = \
+                           parallelMap(updateCPUs,
+                                       lambda f: f.calculateTask(),
+                                       helmholtzFrontiers[helmholtzIndex[0]:helmholtzIndex[0] + helmholtzBatch],
+                                       seedRandom=True)
+                """
+            badIndices = []
+            endingIndex = min(helmholtzIndex[0] + helmholtzBatch, len(helmholtzFrontiers))
+            for i in range(helmholtzIndex[0], endingIndex):
+                helmholtzFrontiers[i].task = newTasks[i - helmholtzIndex[0]]
+                if helmholtzFrontiers[i].task is None: badIndices.append(i)
+            # Permanently kill anything which failed to give a task
+            for i in reversed(badIndices):
+                assert helmholtzFrontiers[i].task is None
+                del helmholtzFrontiers[i]
+
+
+        # We replace each program in the frontier with its likelihoodSummary
+        # This is because calculating likelihood summaries requires juggling types
+        # And type stuff is expensive!
+        frontiers = [self.replaceProgramsWithLikelihoodSummaries(f).normalize()
+                     for f in frontiers]
+
+        eprint("(ID=%d): Training a recognition model from %d frontiers, %d%% Helmholtz, feature extractor %s." % (
+            self.id, len(frontiers), int(helmholtzRatio * 100), self.featureExtractor.__class__.__name__))
+        eprint("(ID=%d): Got %d Helmholtz frontiers - random Helmholtz training? : %s"%(
+            self.id, len(helmholtzFrontiers), len(helmholtzFrontiers) == 0))
+        eprint("(ID=%d): Contextual? %s" % (self.id, str(self.contextual)))
+        eprint("(ID=%d): Bias optimal? %s" % (self.id, str(biasOptimal)))
+        eprint(f"(ID={self.id}): Aux loss? {auxLoss} (n.b. we train a 'auxiliary' classifier anyway - this controls if gradients propagate back to the future extractor)")
+
+        # The number of Helmholtz samples that we generate at once
+        # Should only affect performance and shouldn't affect anything else
+        helmholtzSamples = []
+
+        optimizer = torch.optim.Adam(self.parameters(), lr=lr, eps=1e-3, amsgrad=True)
+        start = time.time()
+        losses, descriptionLengths, realLosses, dreamLosses, realMDL, dreamMDL = [], [], [], [], [], []
+        classificationLosses = []
+        totalGradientSteps = 0
+        epochs = 9999999
+        for i in range(1, epochs + 1):
+            if timeout and time.time() - start > timeout:
+                break
+
+            if totalGradientSteps > steps:
+                break
+
+            if helmholtzRatio < 1.:
+                permutedFrontiers = list(frontiers)
+                random.shuffle(permutedFrontiers)
+            else:
+                permutedFrontiers = [None]
+
+            finishedSteps = False
+            for frontier in permutedFrontiers:
+                # Randomly decide whether to sample from the generative model
+                dreaming = random.random() < helmholtzRatio
+                if dreaming: frontier = getHelmholtz()
+                self.zero_grad()
+                loss, classificationLoss = \
+                        self.frontierBiasOptimal(frontier, auxiliary=auxLoss, vectorized=vectorized) if biasOptimal \
+                        else self.frontierKL(frontier, auxiliary=auxLoss, vectorized=vectorized)
+                if loss is None:
+                    if not dreaming:
+                        eprint("ERROR: Could not extract features during experience replay.")
+                        eprint("Task is:",frontier.task)
+                        eprint("Aborting - we need to be able to extract features of every actual task.")
+                        assert False
+                    else:
+                        continue
+                if is_torch_invalid(loss):
+                    eprint("Invalid real-data loss!")
+                else:
+                    (loss + classificationLoss).backward()
+                    classificationLosses.append(classificationLoss.data.item())
+                    optimizer.step()
+                    totalGradientSteps += 1
+                    losses.append(loss.data.item())
+                    descriptionLengths.append(min(-e.logPrior for e in frontier))
+                    if dreaming:
+                        dreamLosses.append(losses[-1])
+                        dreamMDL.append(descriptionLengths[-1])
+                    else:
+                        realLosses.append(losses[-1])
+                        realMDL.append(descriptionLengths[-1])
+                    if totalGradientSteps > steps:
+                        break # Stop iterating, then print epoch and loss, then break to finish.
+                        
+            if (i == 1 or i % 10 == 0) and losses:
+                eprint("(ID=%d): " % self.id, "Epoch", i, "Loss", mean(losses))
+                if realLosses and dreamLosses:
+                    eprint("(ID=%d): " % self.id, "\t\t(real loss): ", mean(realLosses), "\t(dream loss):", mean(dreamLosses))
+                eprint("(ID=%d): " % self.id, "\tvs MDL (w/o neural net)", mean(descriptionLengths))
+                if realMDL and dreamMDL:
+                    eprint("\t\t(real MDL): ", mean(realMDL), "\t(dream MDL):", mean(dreamMDL))
+                eprint("(ID=%d): " % self.id, "\t%d cumulative gradient steps. %f steps/sec"%(totalGradientSteps,
+                                                                       totalGradientSteps/(time.time() - start)))
+                eprint("(ID=%d): " % self.id, "\t%d-way auxiliary classification loss"%len(self.grammar.primitives),sum(classificationLosses)/len(classificationLosses))
+                losses, descriptionLengths, realLosses, dreamLosses, realMDL, dreamMDL = [], [], [], [], [], []
+                classificationLosses = []
+                gc.collect()
+        
+        eprint("(ID=%d): " % self.id, " Trained recognition model in",time.time() - start,"seconds")
+        self.trained=True
+        return self
+
+    def sampleHelmholtz(self, requests, statusUpdate=None, seed=None):
+        if seed is not None:
+            random.seed(seed)
+        request = random.choice(requests)
+
+        program = self.generativeModel.sample(request, maximumDepth=6, maxAttempts=100)
+        if program is None:
+            return None
+        task = self.featureExtractor.taskOfProgram(program, request)
+
+        if statusUpdate is not None:
+            flushEverything()
+        if task is None:
+            return None
+
+        if hasattr(self.featureExtractor, 'lexicon'):
+            if self.featureExtractor.tokenize(task.examples) is None:
+                return None
+        
+        ll = self.generativeModel.logLikelihood(request, program)
+        frontier = Frontier([FrontierEntry(program=program,
+                                           logLikelihood=0., logPrior=ll)],
+                            task=task)
+        return frontier
+
+    def sampleManyHelmholtz(self, requests, N, CPUs):
+        eprint("Sampling %d programs from the prior on %d CPUs..." % (N, CPUs))
+        flushEverything()
+        frequency = N / 50
+        startingSeed = random.random()
+
+        # Sequentially for ensemble training.
+        samples = [self.sampleHelmholtz(requests,
+                                           statusUpdate='.' if n % frequency == 0 else None,
+                                           seed=startingSeed + n) for n in range(N)]
+
+        # (cathywong) Disabled for ensemble training. 
+        # samples = parallelMap(
+        #     1,
+        #     lambda n: self.sampleHelmholtz(requests,
+        #                                    statusUpdate='.' if n % frequency == 0 else None,
+        #                                    seed=startingSeed + n),
+        #     range(N))
+        eprint()
+        flushEverything()
+        samples = [z for z in samples if z is not None]
+        eprint()
+        eprint("Got %d/%d valid samples." % (len(samples), N))
+        flushEverything()
+
+        return samples
+
+    def enumerateFrontiers(self,
+                           tasks,
+                           enumerationTimeout=None,
+                           testing=False,
+                           solver=None,
+                           CPUs=1,
+                           frontierSize=None,
+                           maximumFrontier=None,
+                           evaluationTimeout=None):
+        with timing("Evaluated recognition model"):
+            grammars = {task: self.grammarOfTask(task)
+                        for task in tasks}
+            #untorch seperately to make sure you filter out None grammars
+            grammars = {task: grammar.untorch() for task, grammar in grammars.items() if grammar is not None}
+
+        return multicoreEnumeration(grammars, tasks,
+                                    testing=testing,
+                                    solver=solver,
+                                    enumerationTimeout=enumerationTimeout,
+                                    CPUs=CPUs, maximumFrontier=maximumFrontier,
+                                    evaluationTimeout=evaluationTimeout)
+
+
+class RecurrentFeatureExtractor(nn.Module):
+    def __init__(self, _=None,
+                 tasks=None,
+                 cuda=False,
+                 # what are the symbols that can occur in the inputs and
+                 # outputs
+                 lexicon=None,
+                 # how many hidden units
+                 H=32,
+                 # Should the recurrent units be bidirectional?
+                 bidirectional=False,
+                 # What should be the timeout for trying to construct Helmholtz tasks?
+                 helmholtzTimeout=0.25,
+                 # What should be the timeout for running a Helmholtz program?
+                 helmholtzEvaluationTimeout=0.01):
+        super(RecurrentFeatureExtractor, self).__init__()
+
+        assert tasks is not None, "You must provide a list of all of the tasks, both those that have been hit and those that have not been hit. Input examples are sampled from these tasks."
+
+        # maps from a requesting type to all of the inputs that we ever saw with that request
+        self.requestToInputs = {
+            tp: [list(map(fst, t.examples)) for t in tasks if t.request == tp ]
+            for tp in {t.request for t in tasks}
+        }
+
+        inputTypes = {t
+                      for task in tasks
+                      for t in task.request.functionArguments()}
+        # maps from a type to all of the inputs that we ever saw having that type
+        self.argumentsWithType = {
+            tp: [ x
+                  for t in tasks
+                  for xs,_ in t.examples
+                  for tpp, x in zip(t.request.functionArguments(), xs)
+                  if tpp == tp]
+            for tp in inputTypes
+        }
+        self.requestToNumberOfExamples = {
+            tp: [ len(t.examples)
+                  for t in tasks if t.request == tp ]
+            for tp in {t.request for t in tasks}
+        }
+        self.helmholtzTimeout = helmholtzTimeout
+        self.helmholtzEvaluationTimeout = helmholtzEvaluationTimeout
+        self.parallelTaskOfProgram = True
+        
+        assert lexicon
+        self.specialSymbols = [
+            "STARTING",  # start of entire sequence
+            "ENDING",  # ending of entire sequence
+            "STARTOFOUTPUT",  # begins the start of the output
+            "ENDOFINPUT"  # delimits the ending of an input - we might have multiple inputs
+        ]
+        lexicon += self.specialSymbols
+        encoder = nn.Embedding(len(lexicon), H)
+        self.encoder = encoder
+
+        self.H = H
+        self.bidirectional = bidirectional
+
+        layers = 1
+
+        model = nn.GRU(H, H, layers, bidirectional=bidirectional)
+        self.model = model
+
+        self.use_cuda = cuda
+        self.lexicon = lexicon
+        self.symbolToIndex = {
+            symbol: index for index,
+            symbol in enumerate(lexicon)}
+        self.startingIndex = self.symbolToIndex["STARTING"]
+        self.endingIndex = self.symbolToIndex["ENDING"]
+        self.startOfOutputIndex = self.symbolToIndex["STARTOFOUTPUT"]
+        self.endOfInputIndex = self.symbolToIndex["ENDOFINPUT"]
+
+        # Maximum number of inputs/outputs we will run the recognition
+        # model on per task
+        # This is an optimization hack
+        self.MAXINPUTS = 100
+
+        if cuda: self.cuda()
+
+    @property
+    def outputDimensionality(self): return self.H
+
+    # modify examples before forward (to turn them into iterables of lexicon)
+    # you should override this if needed
+    def tokenize(self, x): return x
+
+    def symbolEmbeddings(self):
+        return {s: self.encoder(variable([self.symbolToIndex[s]])).squeeze(
+            0).data.cpu().numpy() for s in self.lexicon if not (s in self.specialSymbols)}
+
+    def packExamples(self, examples):
+        """IMPORTANT! xs must be sorted in decreasing order of size because pytorch is stupid"""
+        es = []
+        sizes = []
+        for xs, y in examples:
+            e = [self.startingIndex]
+            for x in xs:
+                for s in x:
+                    e.append(self.symbolToIndex[s])
+                e.append(self.endOfInputIndex)
+            e.append(self.startOfOutputIndex)
+            for s in y:
+                e.append(self.symbolToIndex[s])
+            e.append(self.endingIndex)
+            if es != []:
+                assert len(e) <= len(es[-1]), \
+                    "Examples must be sorted in decreasing order of their tokenized size. This should be transparently handled in recognition.py, so if this assertion fails it isn't your fault as a user of EC but instead is a bug inside of EC."
+            es.append(e)
+            sizes.append(len(e))
+
+        m = max(sizes)
+        # padding
+        for j, e in enumerate(es):
+            es[j] += [self.endingIndex] * (m - len(e))
+
+        x = variable(es, cuda=self.use_cuda)
+        x = self.encoder(x)
+        # x: (batch size, maximum length, E)
+        x = x.permute(1, 0, 2)
+        # x: TxBxE
+        x = pack_padded_sequence(x, sizes)
+        return x, sizes
+
+    def examplesEncoding(self, examples):
+        examples = sorted(examples, key=lambda xs_y: sum(
+            len(z) + 1 for z in xs_y[0]) + len(xs_y[1]), reverse=True)
+        x, sizes = self.packExamples(examples)
+        outputs, hidden = self.model(x)
+        # outputs, sizes = pad_packed_sequence(outputs)
+        # I don't know whether to return the final output or the final hidden
+        # activations...
+        return hidden[0, :, :] + hidden[1, :, :]
+
+    def forward(self, examples):
+        tokenized = self.tokenize(examples)
+        if not tokenized:
+            return None
+
+        if hasattr(self, 'MAXINPUTS') and len(tokenized) > self.MAXINPUTS:
+            tokenized = list(tokenized)
+            random.shuffle(tokenized)
+            tokenized = tokenized[:self.MAXINPUTS]
+        e = self.examplesEncoding(tokenized)
+        # max pool
+        # e,_ = e.max(dim = 0)
+
+        # take the average activations across all of the examples
+        # I think this might be better because we might be testing on data
+        # which has far more o far fewer examples then training
+        e = e.mean(dim=0)
+        return e
+
+    def featuresOfTask(self, t):
+        if hasattr(self, 'useFeatures'):
+            f = self(t.features)
+        else:
+            # Featurize the examples directly.
+            f = self(t.examples)
+        return f
+
+    def taskOfProgram(self, p, tp):
+        # half of the time we randomly mix together inputs
+        # this gives better generalization on held out tasks
+        # the other half of the time we train on sets of inputs in the training data
+        # this gives better generalization on unsolved training tasks
+        if random.random() < 0.5:
+            def randomInput(t): return random.choice(self.argumentsWithType[t])
+            # Loop over the inputs in a random order and pick the first ones that
+            # doesn't generate an exception
+
+            startTime = time.time()
+            examples = []
+            while True:
+                # TIMEOUT! this must not be a very good program
+                if time.time() - startTime > self.helmholtzTimeout: return None
+
+                # Grab some random inputs
+                xs = [randomInput(t) for t in tp.functionArguments()]
+                try:
+                    y = runWithTimeout(lambda: p.runWithArguments(xs), self.helmholtzEvaluationTimeout)
+                    examples.append((tuple(xs),y))
+                    if len(examples) >= random.choice(self.requestToNumberOfExamples[tp]):
+                        return Task("Helmholtz", tp, examples)
+                except: continue
+
+        else:
+            candidateInputs = list(self.requestToInputs[tp])
+            random.shuffle(candidateInputs)
+            for xss in candidateInputs:
+                ys = []
+                for xs in xss:
+                    try: y = runWithTimeout(lambda: p.runWithArguments(xs), self.helmholtzEvaluationTimeout)
+                    except: break
+                    ys.append(y)
+                if len(ys) == len(xss):
+                    return Task("Helmholtz", tp, list(zip(xss, ys)))
+            return None
+                
+            
+    
+class LowRank(nn.Module):
+    """
+    Module that outputs a rank R matrix of size m by n from input of size i.
+    """
+    def __init__(self, i, m, n, r):
+        """
+        i: input dimension
+        m: output rows
+        n: output columns
+        r: maximum rank. if this is None then the output will be full-rank
+        """
+        super(LowRank, self).__init__()
+
+        self.m = m
+        self.n = n
+        
+        maximumPossibleRank = min(m, n)
+        if r is None: r = maximumPossibleRank
+        
+        if r < maximumPossibleRank:
+            self.factored = True
+            self.A = nn.Linear(i, m*r)
+            self.B = nn.Linear(i, n*r)
+            self.r = r
+        else:
+            self.factored = False
+            self.M = nn.Linear(i, m*n)
+
+    def forward(self, x):
+        sz = x.size()
+        if len(sz) == 1:
+            B = 1
+            x = x.unsqueeze(0)
+            needToSqueeze = True
+        elif len(sz) == 2:
+            B = sz[0]
+            needToSqueeze = False
+        else:
+            assert False, "LowRank expects either a 1-dimensional tensor or a 2-dimensional tensor"
+
+        if self.factored:
+            a = self.A(x).view(B, self.m, self.r)
+            b = self.B(x).view(B, self.r, self.n)
+            y = a @ b
+        else:
+            y = self.M(x).view(B, self.m, self.n)
+        if needToSqueeze:
+            y = y.squeeze(0)
+        return y
+            
+            
+            
+
+class DummyFeatureExtractor(nn.Module):
+    def __init__(self, tasks, testingTasks=[], cuda=False):
+        super(DummyFeatureExtractor, self).__init__()
+        self.outputDimensionality = 1
+        self.recomputeTasks = False
+    def featuresOfTask(self, t):
+        return variable([0.]).float()
+    def featuresOfTasks(self, ts):
+        return variable([[0.]]*len(ts)).float()
+    def taskOfProgram(self, p, t):
+        return Task("dummy task", t, [])
+
+class RandomFeatureExtractor(nn.Module):
+    def __init__(self, tasks):
+        super(RandomFeatureExtractor, self).__init__()
+        self.outputDimensionality = 1
+        self.recomputeTasks = False
+    def featuresOfTask(self, t):
+        return variable([random.random()]).float()
+    def featuresOfTasks(self, ts):
+        return variable([[random.random()] for _ in range(len(ts)) ]).float()
+    def taskOfProgram(self, p, t):
+        return Task("dummy task", t, [])
+
+class Flatten(nn.Module):
+    def __init__(self):
+        super(Flatten, self).__init__()
+
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+class ImageFeatureExtractor(nn.Module):
+    def __init__(self, inputImageDimension, resizedDimension=None,
+                 channels=1):
+        super(ImageFeatureExtractor, self).__init__()
+        
+        self.resizedDimension = resizedDimension or inputImageDimension
+        self.inputImageDimension = inputImageDimension
+        self.channels = channels
+
+        def conv_block(in_channels, out_channels):
+            return nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, 3, padding=1),
+                # nn.BatchNorm2d(out_channels),
+                nn.ReLU(),
+                nn.MaxPool2d(2)
+            )
+
+        # channels for hidden
+        hid_dim = 64
+        z_dim = 64
+
+        self.encoder = nn.Sequential(
+            conv_block(channels, hid_dim),
+            conv_block(hid_dim, hid_dim),
+            conv_block(hid_dim, hid_dim),
+            conv_block(hid_dim, z_dim),
+            Flatten()
+        )
+        
+        # Each layer of the encoder halves the dimension, except for the last layer which flattens
+        outputImageDimensionality = self.resizedDimension/(2**(len(self.encoder) - 1))
+        self.outputDimensionality = int(z_dim*outputImageDimensionality*outputImageDimensionality)
+
+    def forward(self, v):
+        """1 channel: v: BxWxW or v:WxW
+        > 1 channel: v: BxCxWxW or v:CxWxW"""
+
+        insertBatch = False
+        variabled = variable(v).float()
+        if self.channels == 1: # insert channel dimension
+            if len(variabled.shape) == 3: # batching
+                variabled = variabled[:,None,:,:]
+            elif len(variabled.shape) == 2: # no batching
+                variabled = variabled[None,:,:]
+                insertBatch = True
+            else: assert False
+        else: # expect to have a channel dimension
+            if len(variabled.shape) == 4:
+                pass
+            elif len(variabled.shape) == 3:
+                insertBatch = True
+            else: assert False                
+
+        if insertBatch: variabled = torch.unsqueeze(variabled, 0)
+        
+        y = self.encoder(variabled)
+        if insertBatch: y = y[0,:]
+        return y
+
+class JSONFeatureExtractor(object):
+    def __init__(self, tasks, cudaFalse):
+        # self.averages, self.deviations = Task.featureMeanAndStandardDeviation(tasks)
+        # self.outputDimensionality = len(self.averages)
+        self.cuda = cuda
+        self.tasks = tasks
+
+    def stringify(self, x):
+        # No whitespace #maybe kill the seperators
+        return json.dumps(x, separators=(',', ':'))
+
+    def featuresOfTask(self, t):
+        # >>> t.request to get the type
+        # >>> t.examples to get input/output examples
+        # this might actually be okay, because the input should just be nothing
+        #return [(self.stringify(inputs), self.stringify(output))
+        #        for (inputs, output) in t.examples]
+        return [(list(output),) for (inputs, output) in t.examples]

dreamcoder/task.py

@@ -0,0 +1,244 @@
+from dreamcoder.program import *
+from dreamcoder.differentiation import *
+
+import signal
+
+
+class EvaluationTimeout(Exception):
+    pass
+
+
+EVALUATIONTABLE = {}
+
+
+class Task(object):
+    def __init__(self, name, request, examples, features=None, cache=False):
+        '''request: the type of this task
+        examples: list of tuples of (input, output). input should be a tuple, with one entry for each argument
+        cache: should program evaluations be cached?
+        features: list of floats.'''
+        self.cache = cache
+        self.features = features
+        self.request = request
+        self.name = name
+        self.examples = examples
+        if len(self.examples) > 0:
+            assert all(len(xs) == len(examples[0][0])
+                       for xs, _ in examples), \
+                "(for task %s) FATAL: Number of arguments varies." % name
+
+    def __str__(self):
+        if self.supervision is None:
+            return self.name
+        else:
+            return self.name + " (%s)"%self.supervision
+
+    def __repr__(self):
+        return "Task(name={self.name}, request={self.request}, examples={self.examples}"\
+            .format(self=self)
+
+    def __eq__(self, o): return self.name == o.name
+
+    def __ne__(self, o): return not (self == o)
+
+    def __hash__(self): return hash(self.name)
+
+    def describe(self):
+        description = ["%s : %s" % (self.name, self.request)]
+        for xs, y in self.examples:
+            if len(xs) == 1:
+                description.append("f(%s) = %s" % (xs[0], y))
+            else:
+                description.append("f%s = %s" % (xs, y))
+        return "\n".join(description)
+
+    def predict(self, f, x):
+        for a in x:
+            f = f(a)
+        return f
+
+    @property
+    def supervision(self):
+        if not hasattr(self, 'supervisedSolution'): return None
+        return self.supervisedSolution
+
+    def check(self, e, timeout=None):
+        if timeout is not None:
+            def timeoutCallBack(_1, _2): raise EvaluationTimeout()
+        try:
+            signal.signal(signal.SIGVTALRM, timeoutCallBack)
+            signal.setitimer(signal.ITIMER_VIRTUAL, timeout)
+
+            try:
+                f = e.evaluate([])
+            except IndexError:
+                # free variable
+                return False
+            except Exception as e:
+                eprint("Exception during evaluation:", e)
+                return False
+
+            for x, y in self.examples:
+                if self.cache and (x, e) in EVALUATIONTABLE:
+                    p = EVALUATIONTABLE[(x, e)]
+                else:
+                    try:
+                        p = self.predict(f, x)
+                    except BaseException:
+                        p = None
+                    if self.cache:
+                        EVALUATIONTABLE[(x, e)] = p
+                if p != y:
+                    if timeout is not None:
+                        signal.signal(signal.SIGVTALRM, lambda *_: None)
+                        signal.setitimer(signal.ITIMER_VIRTUAL, 0)
+                    return False
+
+            return True
+        # except e:
+            # eprint(e)
+            # assert(False)
+        except EvaluationTimeout:
+            eprint("Timed out while evaluating", e)
+            return False
+        finally:
+            if timeout is not None:
+                signal.signal(signal.SIGVTALRM, lambda *_: None)
+                signal.setitimer(signal.ITIMER_VIRTUAL, 0)
+
+    def logLikelihood(self, e, timeout=None):
+        if self.check(e, timeout):
+            return 0.0
+        else:
+            return NEGATIVEINFINITY
+
+    @staticmethod
+    def featureMeanAndStandardDeviation(tasks):
+        dimension = len(tasks[0].features)
+        averages = [sum(t.features[j] for t in tasks) / float(len(tasks))
+                    for j in range(dimension)]
+        variances = [sum((t.features[j] -
+                          averages[j])**2 for t in tasks) /
+                     float(len(tasks)) for j in range(dimension)]
+        standardDeviations = [v**0.5 for v in variances]
+        for j, s in enumerate(standardDeviations):
+            if s == 0.:
+                eprint(
+                    "WARNING: Feature %d is always %f" %
+                    (j + 1, averages[j]))
+        return averages, standardDeviations
+
+    def as_json_dict(self):
+        return {
+            "name": self.name,
+            "request": str(self.request),
+            "examples": [{"inputs": x, "output": y} for x, y in self.examples]
+        }
+
+
+class DifferentiableTask(Task):
+
+    def __init__(self, name, request, examples, _=None,
+                 features=None, BIC=1., loss=None, likelihoodThreshold=None,
+                 steps=50, restarts=300, lr=0.5, decay=0.5, grow=1.2, actualParameters=None,
+                 temperature=1., maxParameters=None, clipLoss=None, clipOutput=None):
+        assert loss is not None
+        self.temperature = temperature
+        self.actualParameters = actualParameters
+        self.maxParameters = maxParameters
+        self.loss = loss
+        self.BIC = BIC
+        self.likelihoodThreshold = likelihoodThreshold
+
+        arguments = {"parameterPenalty": BIC * math.log(len(examples)),
+                     "temperature": temperature,
+                     "steps": steps, "restarts": restarts, "lr": lr, "decay": decay, "grow": grow,
+                     "maxParameters": maxParameters,
+                     "lossThreshold": -likelihoodThreshold}
+        if clipLoss is not None: arguments['clipLoss'] = float(clipLoss)
+        if clipOutput is not None: arguments['clipOutput'] = float(clipOutput)
+        if actualParameters is not None: arguments['actualParameters'] = int(actualParameters)
+        
+        self.specialTask = ("differentiable",
+                            arguments)
+
+        super(
+            DifferentiableTask,
+            self).__init__(
+            name,
+            request,
+            examples,
+            features,
+            cache=False)
+
+    def logLikelihood(self, e, timeout=None):
+        assert timeout is None, "timeout not implemented for differentiable tasks, but not for any good reason."
+        e, parameters = PlaceholderVisitor.execute(e)
+        if self.maxParameters is not None and len(
+                parameters) > self.maxParameters:
+            return NEGATIVEINFINITY
+        if self.actualParameters is not None and len(
+                parameters) > self.actualParameters:
+            return NEGATIVEINFINITY
+        f = e.evaluate([])
+
+        loss = sum(self.loss(self.predict(f, xs), y)
+                   for xs, y in self.examples) / float(len(self.examples))
+        if isinstance(loss, DN):
+            try:
+                loss = loss.restartingOptimize(
+                    parameters,
+                    lr=self.specialTask[1]["lr"],
+                    steps=self.specialTask[1]["steps"],
+                    decay=self.specialTask[1]["decay"],
+                    grow=self.specialTask[1]["grow"],
+                    attempts=self.specialTask[1]["restarts"],
+                    update=None)
+            except InvalidLoss:
+                loss = POSITIVEINFINITY
+
+        # BIC penalty
+        penalty = self.BIC * len(parameters) * math.log(len(self.examples))
+
+        if self.likelihoodThreshold is not None:
+            if loss > -self.likelihoodThreshold:
+                return NEGATIVEINFINITY
+            else:
+                return -penalty
+        else:
+            return -loss / self.temperature - penalty
+
+
+def squaredErrorLoss(prediction, target):
+    d = prediction - target
+    return d * d
+
+
+def l1loss(prediction, target):
+    return abs(prediction - target)
+
+
+class PlaceholderVisitor(object):
+    def __init__(self): self.parameters = []
+
+    def primitive(self, e):
+        if e.name == 'REAL':
+            placeholder = Placeholder.named("REAL_", random.random())
+            self.parameters.append(placeholder)
+            return Primitive(e.name, e.tp, placeholder)
+        return e
+
+    def invented(self, e): return e.body.visit(self)
+
+    def abstraction(self, e): return Abstraction(e.body.visit(self))
+
+    def application(self, e):
+        return Application(e.f.visit(self), e.x.visit(self))
+
+    def index(self, e): return e
+
+    @staticmethod
+    def execute(e):
+        v = PlaceholderVisitor()
+        e = e.visit(v)
+        return e, v.parameters

dreamcoder/taskBatcher.py

@@ -0,0 +1,200 @@
+from dreamcoder.utilities import eprint
+import random
+
+
+class DefaultTaskBatcher:
+        """Iterates through task batches of the specified size. Defaults to all tasks if taskBatchSize is None."""
+
+        def __init__(self):
+                pass
+
+        def getTaskBatch(self, ec_result, tasks, taskBatchSize, currIteration):
+                if taskBatchSize is None:
+                        taskBatchSize = len(tasks)
+                elif taskBatchSize > len(tasks):
+                        eprint("Task batch size is greater than total number of tasks, aborting.")
+                        assert False
+                
+
+                start = (taskBatchSize * currIteration) % len(tasks)
+                end = start + taskBatchSize
+                taskBatch = (tasks + tasks)[start:end] # Handle wraparound.
+                return taskBatch
+
+class RandomTaskBatcher:
+        """Returns a randomly sampled task batch of the specified size. Defaults to all tasks if taskBatchSize is None."""
+
+        def __init__(self):
+                pass
+
+        def getTaskBatch(self, ec_result, tasks, taskBatchSize, currIteration):
+                if taskBatchSize is None:
+                        taskBatchSize = len(tasks)
+                elif taskBatchSize > len(tasks):
+                        eprint("Task batch size is greater than total number of tasks, aborting.")
+                        assert False
+
+                return random.sample(tasks, taskBatchSize)
+
+class RandomShuffleTaskBatcher:
+        """Randomly shuffles the task batch first, and then iterates through task batches of the specified size like DefaultTaskBatcher.
+           Reshuffles across iterations - intended as benchmark comparison to test the task ordering."""
+        def __init__(self, baseSeed=0): self.baseSeed = baseSeed
+
+        def getTaskBatch(self, ec_result, tasks, taskBatchSize, currIteration):
+                if taskBatchSize is None:
+                        taskBatchSize = len(tasks)
+                elif taskBatchSize > len(tasks):
+                        eprint("Task batch size is greater than total number of tasks, aborting.")
+                        assert False
+                
+                # Reshuffles tasks in a fixed way across epochs for reproducibility.
+                currEpoch = int(int(currIteration * taskBatchSize) / int(len(tasks)))
+
+                shuffledTasks = tasks.copy() # Since shuffle works in place.
+                random.Random(self.baseSeed + currEpoch).shuffle(shuffledTasks)
+
+                shuffledTasksWrap = tasks.copy() # Since shuffle works in place.
+                random.Random(self.baseSeed + currEpoch + 1).shuffle(shuffledTasksWrap)
+
+                start = (taskBatchSize * currIteration) % len(shuffledTasks)
+                end = start + taskBatchSize
+                taskBatch = (shuffledTasks + shuffledTasksWrap)[start:end] # Wraparound nicely.
+
+                return list(set(taskBatch))
+
+class UnsolvedTaskBatcher:
+        """At a given epoch, returns only batches of the tasks that have not been solved at least twice"""
+
+        def __init__(self):
+                self.timesSolved = {} # map from task to times that we have solved it
+                self.start = 0
+
+        def getTaskBatch(self, ec_result, tasks, taskBatchSize, currIteration):
+                assert taskBatchSize is None, "This batching strategy does not support batch sizes"
+
+                for t,f in ec_result.allFrontiers.items():
+                        if f.empty:
+                                self.timesSolved[t] = max(0, self.timesSolved.get(t,0))
+                        else:
+                                self.timesSolved[t] = 1 + self.timesSolved.get(t, 0)
+                return [t for t in tasks if self.timesSolved.get(t,0) < 2 ]
+        
+def entropyRandomBatch(ec_result, tasks, taskBatchSize, randomRatio):
+        numRandom = int(randomRatio * taskBatchSize)
+        numEntropy = taskBatchSize - numRandom
+
+        eprint("Selecting top %d tasks from the %d overall tasks given lowest entropy." % (taskBatchSize, len(tasks)))
+        eprint("Will be selecting %d by lowest entropy and %d randomly." %(numEntropy, numRandom))
+        taskGrammarEntropies = ec_result.recognitionModel.taskGrammarEntropies(tasks)
+        sortedEntropies = sorted(taskGrammarEntropies.items(), key=lambda x:x[1])
+
+        entropyBatch = [task for (task, entropy) in sortedEntropies[:numEntropy]]
+        randomBatch = random.sample([task for (task, entropy) in sortedEntropies[numEntropy:]], numRandom)
+        batch = entropyBatch + randomBatch
+
+        return batch
+
+def kNearestNeighbors(ec_result, tasks, k, task):
+        """Finds the k nearest neighbors in the recognition model logProduction space to a given task."""
+        import numpy as np
+        cosDistance = ec_result.recognitionModel.grammarLogProductionDistanceToTask(task, tasks)
+        argSort = np.argsort(-cosDistance)# Want the greatest similarity.
+        topK = argSort[:k]
+        topKTasks = list(np.array(tasks)[topK])
+        return topKTasks
+
+
+class RandomkNNTaskBatcher:
+        """Chooses a random task and finds the (taskBatchSize - 1) nearest neighbors using the recognition model logits."""
+        def __init__(self):
+                pass
+
+        def getTaskBatch(self, ec_result, tasks, taskBatchSize, currIteration):
+                if taskBatchSize is None:
+                        taskBatchSize = len(tasks)
+                elif taskBatchSize > len(tasks):
+                        eprint("Task batch size is greater than total number of tasks, aborting.")
+                        assert False
+
+                if ec_result.recognitionModel is None:
+                        eprint("No recognition model, falling back on random %d" % taskBatchSize)
+                        return random.sample(tasks, taskBatchSize)
+                else:
+                        randomTask = random.choice(tasks)
+                        kNN = kNearestNeighbors(ec_result, tasks, taskBatchSize - 1, randomTask)
+                        return [randomTask] + kNN
+
+class RandomLowEntropykNNTaskBatcher:
+        """Choose a random task from the 10 unsolved with the lowest entropy, and finds the (taskBatchSize - 1) nearest neighbors using the recognition model logits."""
+        def __init__(self):
+                pass
+
+        def getTaskBatch(self, ec_result, tasks, taskBatchSize, currIteration):
+                unsolvedTasks = [t for t in tasks if ec_result.allFrontiers[t].empty]
+
+                if taskBatchSize is None:
+                        return unsolvedTasks
+                elif taskBatchSize > len(tasks):
+                        eprint("Task batch size is greater than total number of tasks, aborting.")
+                        assert False
+
+                if ec_result.recognitionModel is None:
+                        eprint("No recognition model, falling back on random %d tasks from the remaining %d" %(taskBatchSize, len(unsolvedTasks)))
+                        return random.sample(unsolvedTasks, taskBatchSize)
+                else:
+                        lowEntropyUnsolved = entropyRandomBatch(ec_result, unsolvedTasks, taskBatchSize, randomRatio=0)
+                        randomTask = random.choice(lowEntropyUnsolved)
+                        kNN = kNearestNeighbors(ec_result, tasks, taskBatchSize - 1, randomTask)
+                        return [randomTask] + kNN
+
+
+class UnsolvedEntropyTaskBatcher:
+        """Returns tasks that have never been solved at any previous iteration.
+           Given a task batch size, returns the unsolved tasks with the lowest entropy."""
+        def __init__(self):
+                pass
+
+        def getTaskBatch(self, ec_result, tasks, taskBatchSize, currIteration):
+                unsolvedTasks = [t for t in tasks if ec_result.allFrontiers[t].empty]
+
+                if taskBatchSize is None:
+                        return unsolvedTasks
+                elif taskBatchSize > len(tasks):
+                        eprint("Task batch size is greater than total number of tasks, aborting.")
+                        assert False
+
+                if ec_result.recognitionModel is None:
+                        eprint("No recognition model, falling back on random %d tasks from the remaining %d" %(taskBatchSize, len(unsolvedTasks)))
+                        return random.sample(unsolvedTasks, taskBatchSize)
+                else:
+                        return entropyRandomBatch(ec_result, unsolvedTasks, taskBatchSize, randomRatio=0)
+
+class UnsolvedRandomEntropyTaskBatcher:
+        """Returns tasks that have never been solved at any previous iteration.
+           Given a task batch size, returns a mix of unsolved tasks with percentRandom 
+           selected randomly and the remaining selected by lowest entropy.""" 
+        def __init__(self):
+                pass
+
+        def getTaskBatch(self, ec_result, tasks, taskBatchSize, currIteration):
+                unsolvedTasks = [t for t in tasks if ec_result.allFrontiers[t].empty]
+
+                if taskBatchSize is None:
+                        return unsolvedTasks
+                elif taskBatchSize > len(tasks):
+                        eprint("Task batch size is greater than total number of tasks, aborting.")
+                        assert False
+
+                if ec_result.recognitionModel is None:
+                        eprint("No recognition model, falling back on random %d tasks from the remaining %d" %(taskBatchSize, len(unsolvedTasks)))
+                        return random.sample(unsolvedTasks, taskBatchSize)
+                else:
+                        return entropyRandomBatch(ec_result, unsolvedTasks, taskBatchSize, randomRatio=.5)
+
+
+
+
+
+                
+

dreamcoder/type.py

@@ -0,0 +1,378 @@
+class UnificationFailure(Exception):
+    pass
+
+
+class Occurs(UnificationFailure):
+    pass
+
+
+class Type(object):
+    def __str__(self): return self.show(True)
+
+    def __repr__(self): return str(self)
+
+    @staticmethod
+    def fromjson(j):
+        if "index" in j: return TypeVariable(j["index"])
+        if "constructor" in j: return TypeConstructor(j["constructor"],
+                                                      [ Type.fromjson(a) for a in j["arguments"] ])
+        assert False
+
+
+class TypeConstructor(Type):
+    def __init__(self, name, arguments):
+        self.name = name
+        self.arguments = arguments
+        self.isPolymorphic = any(a.isPolymorphic for a in arguments)
+
+    def makeDummyMonomorphic(self, mapping=None):
+        mapping = mapping if mapping is not None else {}
+        return TypeConstructor(self.name,
+                               [ a.makeDummyMonomorphic(mapping) for a in self.arguments ])
+
+    def __eq__(self, other):
+        return isinstance(other, TypeConstructor) and \
+            self.name == other.name and \
+            all(x == y for x, y in zip(self.arguments, other.arguments))
+
+    def __hash__(self): return hash((self.name,) + tuple(self.arguments))
+
+    def __ne__(self, other):
+        return not (self == other)
+
+    def show(self, isReturn):
+        if self.name == ARROW:
+            if isReturn:
+                return "%s %s %s" % (self.arguments[0].show(
+                    False), ARROW, self.arguments[1].show(True))
+            else:
+                return "(%s %s %s)" % (self.arguments[0].show(
+                    False), ARROW, self.arguments[1].show(True))
+        elif self.arguments == []:
+            return self.name
+        else:
+            return "%s(%s)" % (self.name, ", ".join(x.show(True)
+                                                    for x in self.arguments))
+
+    def json(self):
+        return {"constructor": self.name,
+                "arguments": [a.json() for a in self.arguments]}
+
+
+    def isArrow(self): return self.name == ARROW
+
+    def functionArguments(self):
+        if self.name == ARROW:
+            xs = self.arguments[1].functionArguments()
+            return [self.arguments[0]] + xs
+        return []
+
+    def returns(self):
+        if self.name == ARROW:
+            return self.arguments[1].returns()
+        else:
+            return self
+
+    def apply(self, context):
+        if not self.isPolymorphic:
+            return self
+        return TypeConstructor(self.name,
+                               [x.apply(context) for x in self.arguments])
+
+    def applyMutable(self, context):
+        if not self.isPolymorphic:
+            return self
+        return TypeConstructor(self.name,
+                               [x.applyMutable(context) for x in self.arguments])
+
+    def occurs(self, v):
+        if not self.isPolymorphic:
+            return False
+        return any(x.occurs(v) for x in self.arguments)
+
+    def negateVariables(self):
+        return TypeConstructor(self.name,
+                               [a.negateVariables() for a in self.arguments])
+
+    def instantiate(self, context, bindings=None):
+        if not self.isPolymorphic:
+            return context, self
+        if bindings is None:
+            bindings = {}
+        newArguments = []
+        for x in self.arguments:
+            (context, x) = x.instantiate(context, bindings)
+            newArguments.append(x)
+        return (context, TypeConstructor(self.name, newArguments))
+
+    def instantiateMutable(self, context, bindings=None):
+        if not self.isPolymorphic:
+            return self
+        if bindings is None:
+            bindings = {}
+        newArguments = []
+        return TypeConstructor(self.name, [x.instantiateMutable(context, bindings)
+                                           for x in self.arguments ])
+        
+
+    def canonical(self, bindings=None):
+        if not self.isPolymorphic:
+            return self
+        if bindings is None:
+            bindings = {}
+        return TypeConstructor(self.name,
+                               [x.canonical(bindings) for x in self.arguments])
+
+
+class TypeVariable(Type):
+    def __init__(self, j):
+        assert isinstance(j, int)
+        self.v = j
+        self.isPolymorphic = True
+
+    def makeDummyMonomorphic(self, mapping=None):
+        mapping = mapping if mapping is not None else {}
+        if self.v  not in mapping:
+            mapping[self.v] = TypeConstructor(f"dummy_type_{len(mapping)}", [])
+        return mapping[self.v]
+        
+
+    def __eq__(self, other):
+        return isinstance(other, TypeVariable) and self.v == other.v
+
+    def __ne__(self, other): return not (self.v == other.v)
+
+    def __hash__(self): return self.v
+
+    def show(self, _): return "t%d" % self.v
+
+    def json(self):
+        return {"index": self.v}
+
+    def returns(self): return self
+
+    def isArrow(self): return False
+
+    def functionArguments(self): return []
+
+    def apply(self, context):
+        for v, t in context.substitution:
+            if v == self.v:
+                return t.apply(context)
+        return self
+
+    def applyMutable(self, context):
+        s = context.substitution[self.v]
+        if s is None: return self
+        new = s.applyMutable(context)
+        context.substitution[self.v] = new
+        return new
+
+    def occurs(self, v): return v == self.v
+
+    def instantiate(self, context, bindings=None):
+        if bindings is None:
+            bindings = {}
+        if self.v in bindings:
+            return (context, bindings[self.v])
+        new = TypeVariable(context.nextVariable)
+        bindings[self.v] = new
+        context = Context(context.nextVariable + 1, context.substitution)
+        return (context, new)
+
+    def instantiateMutable(self, context, bindings=None):
+        if bindings is None: bindings = {}
+        if self.v in bindings: return bindings[self.v]
+        new = context.makeVariable()
+        bindings[self.v] = new
+        return new
+
+    def canonical(self, bindings=None):
+        if bindings is None:
+            bindings = {}
+        if self.v in bindings:
+            return bindings[self.v]
+        new = TypeVariable(len(bindings))
+        bindings[self.v] = new
+        return new
+
+    def negateVariables(self):
+        return TypeVariable(-1 - self.v)
+
+
+class Context(object):
+    def __init__(self, nextVariable=0, substitution=[]):
+        self.nextVariable = nextVariable
+        self.substitution = substitution
+
+    def extend(self, j, t):
+        return Context(self.nextVariable, [(j, t)] + self.substitution)
+
+    def makeVariable(self):
+        return (Context(self.nextVariable + 1, self.substitution),
+                TypeVariable(self.nextVariable))
+
+    def unify(self, t1, t2):
+        t1 = t1.apply(self)
+        t2 = t2.apply(self)
+        if t1 == t2:
+            return self
+        # t1&t2 are not equal
+        if not t1.isPolymorphic and not t2.isPolymorphic:
+            raise UnificationFailure(t1, t2)
+
+        if isinstance(t1, TypeVariable):
+            if t2.occurs(t1.v):
+                raise Occurs()
+            return self.extend(t1.v, t2)
+        if isinstance(t2, TypeVariable):
+            if t1.occurs(t2.v):
+                raise Occurs()
+            return self.extend(t2.v, t1)
+        if t1.name != t2.name:
+            raise UnificationFailure(t1, t2)
+        k = self
+        for x, y in zip(t2.arguments, t1.arguments):
+            k = k.unify(x, y)
+        return k
+
+    def __str__(self):
+        return "Context(next = %d, {%s})" % (self.nextVariable, ", ".join(
+            "t%d ||> %s" % (k, v.apply(self)) for k, v in self.substitution))
+
+    def __repr__(self): return str(self)
+
+class MutableContext(object):
+    def __init__(self):
+        self.substitution = []
+
+    def extend(self,i,t):
+        assert self.substitution[i] is None
+        self.substitution[i] = t
+
+    def makeVariable(self):
+        self.substitution.append(None)
+        return TypeVariable(len(self.substitution) - 1)
+
+    def unify(self, t1, t2):
+        t1 = t1.applyMutable(self)
+        t2 = t2.applyMutable(self)
+
+        if t1 == t2: return
+
+        # t1&t2 are not equal
+        if not t1.isPolymorphic and not t2.isPolymorphic:
+            raise UnificationFailure(t1, t2)
+
+        if isinstance(t1, TypeVariable):
+            if t2.occurs(t1.v):
+                raise Occurs()
+            self.extend(t1.v, t2)
+            return 
+        if isinstance(t2, TypeVariable):
+            if t1.occurs(t2.v):
+                raise Occurs()
+            self.extend(t2.v, t1)
+            return 
+        if t1.name != t2.name:
+            raise UnificationFailure(t1, t2)
+        
+        for x, y in zip(t2.arguments, t1.arguments):
+            self.unify(x, y)
+
+
+Context.EMPTY = Context(0, [])
+
+
+def canonicalTypes(ts):
+    bindings = {}
+    return [t.canonical(bindings) for t in ts]
+
+
+def instantiateTypes(context, ts):
+    bindings = {}
+    newTypes = []
+    for t in ts:
+        context, t = t.instantiate(context, bindings)
+        newTypes.append(t)
+    return context, newTypes
+
+
+def baseType(n): return TypeConstructor(n, [])
+
+
+tint = baseType("int")
+treal = baseType("real")
+tbool = baseType("bool")
+tboolean = tbool  # alias
+tcharacter = baseType("char")
+
+
+def tlist(t): return TypeConstructor("list", [t])
+
+
+def tpair(a, b): return TypeConstructor("pair", [a, b])
+
+
+def tmaybe(t): return TypeConstructor("maybe", [t])
+
+
+tstr = tlist(tcharacter)
+t0 = TypeVariable(0)
+t1 = TypeVariable(1)
+t2 = TypeVariable(2)
+
+# regex types
+tpregex = baseType("pregex")
+
+ARROW = "->"
+
+
+def arrow(*arguments):
+    if len(arguments) == 1:
+        return arguments[0]
+    return TypeConstructor(ARROW, [arguments[0], arrow(*arguments[1:])])
+
+
+def inferArg(tp, tcaller):
+    ctx, tp = tp.instantiate(Context.EMPTY)
+    ctx, tcaller = tcaller.instantiate(ctx)
+    ctx, targ = ctx.makeVariable()
+    ctx = ctx.unify(tcaller, arrow(targ, tp))
+    return targ.apply(ctx)
+
+
+def guess_type(xs):
+    """
+    Return a TypeConstructor corresponding to x's python type.
+    Raises an exception if the type cannot be guessed.
+    """
+    if all(isinstance(x, bool) for x in xs):
+        return tbool
+    elif all(isinstance(x, int) for x in xs):
+        return tint
+    elif all(isinstance(x, str) for x in xs):
+        return tstr
+    elif all(isinstance(x, list) for x in xs):
+        return tlist(guess_type([y for ys in xs for y in ys]))
+    else:
+        raise ValueError("cannot guess type from {}".format(xs))
+
+
+def guess_arrow_type(examples):
+    a = len(examples[0][0])
+    input_types = []
+    for n in range(a):
+        input_types.append(guess_type([xs[n] for xs, _ in examples]))
+    output_type = guess_type([y for _, y in examples])
+    return arrow(*(input_types + [output_type]))
+
+def canUnify(t1, t2):
+    k = MutableContext()
+    t1 = t1.instantiateMutable(k)
+    t2 = t2.instantiateMutable(k)
+    try:
+        k.unify(t1, t2)
+        return True
+    except UnificationFailure: return False
+