We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Calls to "main"
In merging #550 I find I need ``ks::main``.  
Looking at this comment in 
https://github.com/microsoft/knossos-ksc/blob/18f654086f349cceb59c30dcbc8de2c534dcb5ec/src/ksc/Cgen.hs#L801

I wonder is it time to look at removing that?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/python/ksc/utils.py`
Content:
```
1 from dataclasses import dataclass
2 from typing import Tuple, Union
3 from collections import namedtuple
4 import itertools
5 
6 import importlib.util
7 import os
8 import numpy as np
9 import subprocess
10 import sysconfig
11 import sys
12 from tempfile import NamedTemporaryFile
13 
14 from ksc.type import Type, tangent_type, make_tuple_if_many
15 
16 class KRecord:
17     """
18     A smoother namedtuple -- like https://pythonhosted.org/pyrecord but using the existing class syntax.
19     Like a 3.7 dataclass, but don't need to decorate each derived class
20 
21     Derive a class from KRecord, declare its fields, and use keyword args in __init__
22 
23     def MyClass(KRecord):
24         cost: float
25         names: List[String]
26 
27         def __init__(cost, names):
28             super().__init__(cost=cost, names=names)
29 
30     And now you have a nice little record class.
31 
32     Construct a MyClass:
33         a = MyClass(1.3, ["fred", "conor", "una"])
34 
35     Compare two MyClasses
36         if a == b: ...
37     
38     Etc
39     """
40 
41     def __init__(self, **args):
42         for (nt,v) in args.items():
43             # assert nt in self.__annotations__  # <- This check will fail for chains of derived classes -- only the deepest has __annotations__ ready yet.
44             setattr(self, nt, v)
45 
46     def __eq__(self, that):
47         if type(self) != type(that):
48             return False
49 
50         for nt in self.__annotations__:
51             if getattr(self, nt) != getattr(that,nt):
52                 return False
53         return True
54 
55 
56 
57 
58 def ensure_list_of_lists(l):
59     """return input, wrapped in a singleton list if its first element is not a list
60 
61        ensure_list_of_lists([])    = []
62        ensure_list_of_lists([1])   = [[1]]
63        ensure_list_of_lists([[1]]) = [[1]]
64        ensure_list_of_lists([[1,2]])        = [[1, 2]]
65        ensure_list_of_lists([[1,2], [3,4]]) = [[1, 2], [3, 4]]
66     """
67 
68     if not isinstance(l, list):
69         raise ValueError("Expect a list")
70     if len(l) < 1:  # Empty list is empty list
71         return l
72     if not isinstance(l[0], list):
73         return [l]
74     else:
75         return l
76 
77 def paren(s):
78     return "(" + s + ")"
79 
80 PYTHON_MODULE_NAME = "ks_mod"
81 
82 def import_module_from_path(module_name, path):
83     # These three lines are for loading a module from a file in Python 3.5+
84     # https://bugs.python.org/issue21436
85     spec = importlib.util.spec_from_file_location(module_name, path)
86     py_out = importlib.util.module_from_spec(spec)
87     spec.loader.exec_module(py_out)
88     return py_out
89 
90 def translate_and_import(source_file_name, *args):
91     from ksc.translate import translate
92     py_out = translate(*args, source_file_name, with_main=False)
93     with NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
94         f.write(f"# AUTOGEN from {source_file_name} via ksc.utils.translate_and_import")
95         f.write(py_out)
96 
97     print(f.name)
98     return import_module_from_path(PYTHON_MODULE_NAME, f.name)
99 
100 def subprocess_run(cmd, env=None):
101     return subprocess.run(cmd, stdout=subprocess.PIPE, env=env).stdout.decode().strip("\n")
102 
103 def get_ksc_paths():
104     if "KSC_RUNTIME_DIR" in os.environ:
105         ksc_runtime_dir = os.environ["KSC_RUNTIME_DIR"]
106     else:
107         ksc_runtime_dir = "./src/runtime"
108 
109     if "KSC_PATH" in os.environ:
110         ksc_path = os.environ["KSC_PATH"]
111     else:
112         ksc_path = "./build/bin/ksc"
113     
114     return ksc_path,ksc_runtime_dir
115 
116 
117 def generate_cpp_from_ks(ks_str, generate_derivatives = False, use_aten = False):
118     ksc_path,_ksc_runtime_dir = get_ksc_paths()
119 
120     with NamedTemporaryFile(mode="w", suffix=".ks", delete=False) as fks:
121         fks.write(ks_str)
122     try:
123         with NamedTemporaryFile(mode="w", suffix=".kso", delete=False) as fkso:
124             with NamedTemporaryFile(mode="w", suffix=".cpp", delete=False) as fcpp:
125                 print("generate_cpp_from_ks:", ksc_path, fks.name)
126                 e = subprocess.run([
127                     ksc_path,
128                     "--generate-cpp" if generate_derivatives else "--generate-cpp-without-diffs",
129                     "--ks-source-file", "src/runtime/prelude.ks",
130                     *(("--ks-source-file", "src/runtime/prelude-aten.ks") if use_aten else ()),
131                     "--ks-source-file", fks.name,
132                     "--ks-output-file", fkso.name,
133                     "--cpp-output-file", fcpp.name
134                 ], capture_output=True, check=True)
135                 print(e.stdout.decode('ascii'))
136                 print(e.stderr.decode('ascii'))
137     except subprocess.CalledProcessError as e:
138         print(f"files {fks.name} {fkso.name} {fcpp.name}")
139         print(f"ks_str=\n{ks_str}")
140         print(e.output.decode('ascii'))
141         print(e.stderr.decode('ascii'))
142         raise
143 
144 
145     # Read from CPP back to string
146     with open(fcpp.name) as f:
147         out = f.read()
148 
149     # only delete these file if no error
150     os.unlink(fks.name)
151     os.unlink(fcpp.name)
152     os.unlink(fkso.name)
153 
154     return out
155 
156 def build_py_module_from_cpp(cpp_str, pybind11_path, use_aten=False):
157     _ksc_path,ksc_runtime_dir = get_ksc_paths()
158 
159     with NamedTemporaryFile(mode="w", suffix=".cpp", delete=False) as fcpp:
160         fcpp.write(cpp_str)
161 
162     extension_suffix = sysconfig.get_config_var('EXT_SUFFIX')
163     if extension_suffix is None:
164         extension_suffix = sysconfig.get_config_var('SO')
165 
166     with NamedTemporaryFile(mode="w", suffix=extension_suffix, delete=False) as fpymod:
167         pass
168     module_path = fpymod.name
169     module_name = os.path.basename(module_path).split(".")[0]
170     python_includes = subprocess_run(
171         [sys.executable, "-m", "pybind11", "--includes"],
172         env={"PYTHONPATH": "extern/pybind11"}
173     )
174     try:
175         cmd = (f"g++ -I{ksc_runtime_dir} -I{pybind11_path}/include "
176                + python_includes
177                + " -Wall -Wno-unused-variable -Wno-unused-but-set-variable"
178                  " -fmax-errors=1"
179                  " -std=c++17"
180                  " -O3"
181                  " -fPIC"
182                  " -shared"
183                + (" -DKS_INCLUDE_ATEN" if use_aten else "")
184                + f" -DPYTHON_MODULE_NAME={module_name}"
185                  f" -o {module_path} "
186                + fcpp.name)
187         print(cmd)
188         subprocess.run(cmd, shell=True, capture_output=True, check=True)
189     except subprocess.CalledProcessError as e:
190         print(f"cpp_file={fcpp.name}")
191         print(cmd)
192         print(e.output.decode('utf-8'))
193         print(e.stderr.decode('utf-8'))
194 
195         raise
196     
197     os.unlink(fcpp.name)
198     return module_name, module_path
199 
200 def mangleType(ty):
201     return ty.shortstr()
202 
203 def mangleTypes(tys):
204     return "".join(mangleType(ty) for ty in tys)
205 
206 
207 def encode_name(s : str) -> str:
208     # TODO: this could be faster
209     return s.\
210         replace('@',"$a").\
211         replace(',',"$_").\
212         replace('.',"$o").\
213         replace('[',"$6").\
214         replace(']',"$9").\
215         replace('<',"$d").\
216         replace('>',"$b").\
217         replace('*',"$x").\
218         replace(':',"$8")
219 
220 def generate_and_compile_cpp_from_ks(ks_str, name_to_call, arg_types, return_type=None, generate_derivatives=False, use_aten=False, pybind11_path="extern/pybind11"):
221 
222     generated_cpp_source = generate_cpp_from_ks(ks_str, generate_derivatives=generate_derivatives, use_aten=use_aten)
223 
224     cpp_str = """
225 #include "knossos-pybind.h"
226 
227 """ + generated_cpp_source + """
228 
229 int ks::main(ks::allocator *) { return 0; };
230 
231 """
232 
233     args_str = mangleTypes(arg_types)
234     name_str = encode_name(f"{name_to_call}@{args_str}")
235     declarations = f"""
236      m.def("entry", with_ks_allocator(&ks::{name_str}));
237     """
238 
239     if generate_derivatives:
240         darg_types = [tangent_type(ty) for ty in arg_types]
241         args_tuple_str = mangleType(make_tuple_if_many(arg_types))
242         dargs_tuple_str = mangleType(make_tuple_if_many(darg_types))
243         dreturn_type_str = mangleType(tangent_type(return_type))
244 
245         fwd_name = encode_name(f"fwd${name_to_call}@{args_tuple_str}")
246         declarations += f"""
247           m.def("fwd_entry", with_ks_allocator(&ks::{fwd_name}));
248         """
249 
250         rev_name = encode_name(f"rev${name_to_call}@{args_tuple_str}")
251         declarations += f"""
252           m.def("rev_entry", with_ks_allocator(&ks::{rev_name}));
253         """
254 
255     cpp_str += """
256 PYBIND11_MODULE(PYTHON_MODULE_NAME, m) {
257     m.def("reset_allocator", []{g_alloc.reset();});
258     m.def("allocator_top", []{ return g_alloc.mark();});
259     m.def("allocator_peak", []{ return g_alloc.peak();});
260 
261     declare_tensor_2<double>(m, "Tensor_2_Float");
262     declare_tensor_2<int>(m, "Tensor_2_Integer");
263 
264 """ + declarations + """
265 }
266 """
267     pybindcpppath = "obj/pybind.cpp"
268     print(f"Saving to {pybindcpppath}")    
269     os.makedirs(os.path.dirname(pybindcpppath), exist_ok=True)
270     with open(pybindcpppath, "w") as fcpp:
271         fcpp.write(cpp_str)
272 
273     module_name, module_path = build_py_module_from_cpp(cpp_str, pybind11_path, use_aten=use_aten)
274     return import_module_from_path(module_name, module_path)
275 
276 def ndgrid_inds(sz):
277     """
278     Return a sequnce of tuples of indices as if generated by nested comprehensions.
279     Example:
280         ndgrid_inds((ni,nj))
281     Returns the same sequence as
282         [(i,j) for i in range(ni) for j in range(nj)]
283 
284     The iterates are always tuples so
285         ndgrid_inds(4)
286     returns
287         [(0,), (1,), (2,), (3,)] 
288 
289     """
290 
291     return itertools.product(*map(range, sz))
292 
293
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/python/ksc/utils.py b/src/python/ksc/utils.py
--- a/src/python/ksc/utils.py
+++ b/src/python/ksc/utils.py
@@ -221,14 +221,11 @@
 
     generated_cpp_source = generate_cpp_from_ks(ks_str, generate_derivatives=generate_derivatives, use_aten=use_aten)
 
-    cpp_str = """
-#include "knossos-pybind.h"
+    cpp_str = f"""
+    #include "knossos-pybind.h"
+    {generated_cpp_source}
 
-""" + generated_cpp_source + """
-
-int ks::main(ks::allocator *) { return 0; };
-
-"""
+    """
 
     args_str = mangleTypes(arg_types)
     name_str = encode_name(f"{name_to_call}@{args_str}")

{"golden_diff": "diff --git a/src/python/ksc/utils.py b/src/python/ksc/utils.py\n--- a/src/python/ksc/utils.py\n+++ b/src/python/ksc/utils.py\n@@ -221,14 +221,11 @@\n \n     generated_cpp_source = generate_cpp_from_ks(ks_str, generate_derivatives=generate_derivatives, use_aten=use_aten)\n \n-    cpp_str = \"\"\"\n-#include \"knossos-pybind.h\"\n+    cpp_str = f\"\"\"\n+    #include \"knossos-pybind.h\"\n+    {generated_cpp_source}\n \n-\"\"\" + generated_cpp_source + \"\"\"\n-\n-int ks::main(ks::allocator *) { return 0; };\n-\n-\"\"\"\n+    \"\"\"\n \n     args_str = mangleTypes(arg_types)\n     name_str = encode_name(f\"{name_to_call}@{args_str}\")\n", "issue": "Calls to \"main\"\nIn merging #550 I find I need ``ks::main``.  \r\nLooking at this comment in \r\nhttps://github.com/microsoft/knossos-ksc/blob/18f654086f349cceb59c30dcbc8de2c534dcb5ec/src/ksc/Cgen.hs#L801\r\n\r\nI wonder is it time to look at removing that?\n", "before_files": [{"content": "from dataclasses import dataclass\nfrom typing import Tuple, Union\nfrom collections import namedtuple\nimport itertools\n\nimport importlib.util\nimport os\nimport numpy as np\nimport subprocess\nimport sysconfig\nimport sys\nfrom tempfile import NamedTemporaryFile\n\nfrom ksc.type import Type, tangent_type, make_tuple_if_many\n\nclass KRecord:\n    \"\"\"\n    A smoother namedtuple -- like https://pythonhosted.org/pyrecord but using the existing class syntax.\n    Like a 3.7 dataclass, but don't need to decorate each derived class\n\n    Derive a class from KRecord, declare its fields, and use keyword args in __init__\n\n    def MyClass(KRecord):\n        cost: float\n        names: List[String]\n\n        def __init__(cost, names):\n            super().__init__(cost=cost, names=names)\n\n    And now you have a nice little record class.\n\n    Construct a MyClass:\n        a = MyClass(1.3, [\"fred\", \"conor\", \"una\"])\n\n    Compare two MyClasses\n        if a == b: ...\n    \n    Etc\n    \"\"\"\n\n    def __init__(self, **args):\n        for (nt,v) in args.items():\n            # assert nt in self.__annotations__  # <- This check will fail for chains of derived classes -- only the deepest has __annotations__ ready yet.\n            setattr(self, nt, v)\n\n    def __eq__(self, that):\n        if type(self) != type(that):\n            return False\n\n        for nt in self.__annotations__:\n            if getattr(self, nt) != getattr(that,nt):\n                return False\n        return True\n\n\n\n\ndef ensure_list_of_lists(l):\n    \"\"\"return input, wrapped in a singleton list if its first element is not a list\n\n       ensure_list_of_lists([])    = []\n       ensure_list_of_lists([1])   = [[1]]\n       ensure_list_of_lists([[1]]) = [[1]]\n       ensure_list_of_lists([[1,2]])        = [[1, 2]]\n       ensure_list_of_lists([[1,2], [3,4]]) = [[1, 2], [3, 4]]\n    \"\"\"\n\n    if not isinstance(l, list):\n        raise ValueError(\"Expect a list\")\n    if len(l) < 1:  # Empty list is empty list\n        return l\n    if not isinstance(l[0], list):\n        return [l]\n    else:\n        return l\n\ndef paren(s):\n    return \"(\" + s + \")\"\n\nPYTHON_MODULE_NAME = \"ks_mod\"\n\ndef import_module_from_path(module_name, path):\n    # These three lines are for loading a module from a file in Python 3.5+\n    # https://bugs.python.org/issue21436\n    spec = importlib.util.spec_from_file_location(module_name, path)\n    py_out = importlib.util.module_from_spec(spec)\n    spec.loader.exec_module(py_out)\n    return py_out\n\ndef translate_and_import(source_file_name, *args):\n    from ksc.translate import translate\n    py_out = translate(*args, source_file_name, with_main=False)\n    with NamedTemporaryFile(mode=\"w\", suffix=\".py\", delete=False) as f:\n        f.write(f\"# AUTOGEN from {source_file_name} via ksc.utils.translate_and_import\")\n        f.write(py_out)\n\n    print(f.name)\n    return import_module_from_path(PYTHON_MODULE_NAME, f.name)\n\ndef subprocess_run(cmd, env=None):\n    return subprocess.run(cmd, stdout=subprocess.PIPE, env=env).stdout.decode().strip(\"\\n\")\n\ndef get_ksc_paths():\n    if \"KSC_RUNTIME_DIR\" in os.environ:\n        ksc_runtime_dir = os.environ[\"KSC_RUNTIME_DIR\"]\n    else:\n        ksc_runtime_dir = \"./src/runtime\"\n\n    if \"KSC_PATH\" in os.environ:\n        ksc_path = os.environ[\"KSC_PATH\"]\n    else:\n        ksc_path = \"./build/bin/ksc\"\n    \n    return ksc_path,ksc_runtime_dir\n\n\ndef generate_cpp_from_ks(ks_str, generate_derivatives = False, use_aten = False):\n    ksc_path,_ksc_runtime_dir = get_ksc_paths()\n\n    with NamedTemporaryFile(mode=\"w\", suffix=\".ks\", delete=False) as fks:\n        fks.write(ks_str)\n    try:\n        with NamedTemporaryFile(mode=\"w\", suffix=\".kso\", delete=False) as fkso:\n            with NamedTemporaryFile(mode=\"w\", suffix=\".cpp\", delete=False) as fcpp:\n                print(\"generate_cpp_from_ks:\", ksc_path, fks.name)\n                e = subprocess.run([\n                    ksc_path,\n                    \"--generate-cpp\" if generate_derivatives else \"--generate-cpp-without-diffs\",\n                    \"--ks-source-file\", \"src/runtime/prelude.ks\",\n                    *((\"--ks-source-file\", \"src/runtime/prelude-aten.ks\") if use_aten else ()),\n                    \"--ks-source-file\", fks.name,\n                    \"--ks-output-file\", fkso.name,\n                    \"--cpp-output-file\", fcpp.name\n                ], capture_output=True, check=True)\n                print(e.stdout.decode('ascii'))\n                print(e.stderr.decode('ascii'))\n    except subprocess.CalledProcessError as e:\n        print(f\"files {fks.name} {fkso.name} {fcpp.name}\")\n        print(f\"ks_str=\\n{ks_str}\")\n        print(e.output.decode('ascii'))\n        print(e.stderr.decode('ascii'))\n        raise\n\n\n    # Read from CPP back to string\n    with open(fcpp.name) as f:\n        out = f.read()\n\n    # only delete these file if no error\n    os.unlink(fks.name)\n    os.unlink(fcpp.name)\n    os.unlink(fkso.name)\n\n    return out\n\ndef build_py_module_from_cpp(cpp_str, pybind11_path, use_aten=False):\n    _ksc_path,ksc_runtime_dir = get_ksc_paths()\n\n    with NamedTemporaryFile(mode=\"w\", suffix=\".cpp\", delete=False) as fcpp:\n        fcpp.write(cpp_str)\n\n    extension_suffix = sysconfig.get_config_var('EXT_SUFFIX')\n    if extension_suffix is None:\n        extension_suffix = sysconfig.get_config_var('SO')\n\n    with NamedTemporaryFile(mode=\"w\", suffix=extension_suffix, delete=False) as fpymod:\n        pass\n    module_path = fpymod.name\n    module_name = os.path.basename(module_path).split(\".\")[0]\n    python_includes = subprocess_run(\n        [sys.executable, \"-m\", \"pybind11\", \"--includes\"],\n        env={\"PYTHONPATH\": \"extern/pybind11\"}\n    )\n    try:\n        cmd = (f\"g++ -I{ksc_runtime_dir} -I{pybind11_path}/include \"\n               + python_includes\n               + \" -Wall -Wno-unused-variable -Wno-unused-but-set-variable\"\n                 \" -fmax-errors=1\"\n                 \" -std=c++17\"\n                 \" -O3\"\n                 \" -fPIC\"\n                 \" -shared\"\n               + (\" -DKS_INCLUDE_ATEN\" if use_aten else \"\")\n               + f\" -DPYTHON_MODULE_NAME={module_name}\"\n                 f\" -o {module_path} \"\n               + fcpp.name)\n        print(cmd)\n        subprocess.run(cmd, shell=True, capture_output=True, check=True)\n    except subprocess.CalledProcessError as e:\n        print(f\"cpp_file={fcpp.name}\")\n        print(cmd)\n        print(e.output.decode('utf-8'))\n        print(e.stderr.decode('utf-8'))\n\n        raise\n    \n    os.unlink(fcpp.name)\n    return module_name, module_path\n\ndef mangleType(ty):\n    return ty.shortstr()\n\ndef mangleTypes(tys):\n    return \"\".join(mangleType(ty) for ty in tys)\n\n\ndef encode_name(s : str) -> str:\n    # TODO: this could be faster\n    return s.\\\n        replace('@',\"$a\").\\\n        replace(',',\"$_\").\\\n        replace('.',\"$o\").\\\n        replace('[',\"$6\").\\\n        replace(']',\"$9\").\\\n        replace('<',\"$d\").\\\n        replace('>',\"$b\").\\\n        replace('*',\"$x\").\\\n        replace(':',\"$8\")\n\ndef generate_and_compile_cpp_from_ks(ks_str, name_to_call, arg_types, return_type=None, generate_derivatives=False, use_aten=False, pybind11_path=\"extern/pybind11\"):\n\n    generated_cpp_source = generate_cpp_from_ks(ks_str, generate_derivatives=generate_derivatives, use_aten=use_aten)\n\n    cpp_str = \"\"\"\n#include \"knossos-pybind.h\"\n\n\"\"\" + generated_cpp_source + \"\"\"\n\nint ks::main(ks::allocator *) { return 0; };\n\n\"\"\"\n\n    args_str = mangleTypes(arg_types)\n    name_str = encode_name(f\"{name_to_call}@{args_str}\")\n    declarations = f\"\"\"\n     m.def(\"entry\", with_ks_allocator(&ks::{name_str}));\n    \"\"\"\n\n    if generate_derivatives:\n        darg_types = [tangent_type(ty) for ty in arg_types]\n        args_tuple_str = mangleType(make_tuple_if_many(arg_types))\n        dargs_tuple_str = mangleType(make_tuple_if_many(darg_types))\n        dreturn_type_str = mangleType(tangent_type(return_type))\n\n        fwd_name = encode_name(f\"fwd${name_to_call}@{args_tuple_str}\")\n        declarations += f\"\"\"\n          m.def(\"fwd_entry\", with_ks_allocator(&ks::{fwd_name}));\n        \"\"\"\n\n        rev_name = encode_name(f\"rev${name_to_call}@{args_tuple_str}\")\n        declarations += f\"\"\"\n          m.def(\"rev_entry\", with_ks_allocator(&ks::{rev_name}));\n        \"\"\"\n\n    cpp_str += \"\"\"\nPYBIND11_MODULE(PYTHON_MODULE_NAME, m) {\n    m.def(\"reset_allocator\", []{g_alloc.reset();});\n    m.def(\"allocator_top\", []{ return g_alloc.mark();});\n    m.def(\"allocator_peak\", []{ return g_alloc.peak();});\n\n    declare_tensor_2<double>(m, \"Tensor_2_Float\");\n    declare_tensor_2<int>(m, \"Tensor_2_Integer\");\n\n\"\"\" + declarations + \"\"\"\n}\n\"\"\"\n    pybindcpppath = \"obj/pybind.cpp\"\n    print(f\"Saving to {pybindcpppath}\")    \n    os.makedirs(os.path.dirname(pybindcpppath), exist_ok=True)\n    with open(pybindcpppath, \"w\") as fcpp:\n        fcpp.write(cpp_str)\n\n    module_name, module_path = build_py_module_from_cpp(cpp_str, pybind11_path, use_aten=use_aten)\n    return import_module_from_path(module_name, module_path)\n\ndef ndgrid_inds(sz):\n    \"\"\"\n    Return a sequnce of tuples of indices as if generated by nested comprehensions.\n    Example:\n        ndgrid_inds((ni,nj))\n    Returns the same sequence as\n        [(i,j) for i in range(ni) for j in range(nj)]\n\n    The iterates are always tuples so\n        ndgrid_inds(4)\n    returns\n        [(0,), (1,), (2,), (3,)] \n\n    \"\"\"\n\n    return itertools.product(*map(range, sz))\n\n", "path": "src/python/ksc/utils.py"}], "after_files": [{"content": "from dataclasses import dataclass\nfrom typing import Tuple, Union\nfrom collections import namedtuple\nimport itertools\n\nimport importlib.util\nimport os\nimport numpy as np\nimport subprocess\nimport sysconfig\nimport sys\nfrom tempfile import NamedTemporaryFile\n\nfrom ksc.type import Type, tangent_type, make_tuple_if_many\n\nclass KRecord:\n    \"\"\"\n    A smoother namedtuple -- like https://pythonhosted.org/pyrecord but using the existing class syntax.\n    Like a 3.7 dataclass, but don't need to decorate each derived class\n\n    Derive a class from KRecord, declare its fields, and use keyword args in __init__\n\n    def MyClass(KRecord):\n        cost: float\n        names: List[String]\n\n        def __init__(cost, names):\n            super().__init__(cost=cost, names=names)\n\n    And now you have a nice little record class.\n\n    Construct a MyClass:\n        a = MyClass(1.3, [\"fred\", \"conor\", \"una\"])\n\n    Compare two MyClasses\n        if a == b: ...\n    \n    Etc\n    \"\"\"\n\n    def __init__(self, **args):\n        for (nt,v) in args.items():\n            # assert nt in self.__annotations__  # <- This check will fail for chains of derived classes -- only the deepest has __annotations__ ready yet.\n            setattr(self, nt, v)\n\n    def __eq__(self, that):\n        if type(self) != type(that):\n            return False\n\n        for nt in self.__annotations__:\n            if getattr(self, nt) != getattr(that,nt):\n                return False\n        return True\n\n\n\n\ndef ensure_list_of_lists(l):\n    \"\"\"return input, wrapped in a singleton list if its first element is not a list\n\n       ensure_list_of_lists([])    = []\n       ensure_list_of_lists([1])   = [[1]]\n       ensure_list_of_lists([[1]]) = [[1]]\n       ensure_list_of_lists([[1,2]])        = [[1, 2]]\n       ensure_list_of_lists([[1,2], [3,4]]) = [[1, 2], [3, 4]]\n    \"\"\"\n\n    if not isinstance(l, list):\n        raise ValueError(\"Expect a list\")\n    if len(l) < 1:  # Empty list is empty list\n        return l\n    if not isinstance(l[0], list):\n        return [l]\n    else:\n        return l\n\ndef paren(s):\n    return \"(\" + s + \")\"\n\nPYTHON_MODULE_NAME = \"ks_mod\"\n\ndef import_module_from_path(module_name, path):\n    # These three lines are for loading a module from a file in Python 3.5+\n    # https://bugs.python.org/issue21436\n    spec = importlib.util.spec_from_file_location(module_name, path)\n    py_out = importlib.util.module_from_spec(spec)\n    spec.loader.exec_module(py_out)\n    return py_out\n\ndef translate_and_import(source_file_name, *args):\n    from ksc.translate import translate\n    py_out = translate(*args, source_file_name, with_main=False)\n    with NamedTemporaryFile(mode=\"w\", suffix=\".py\", delete=False) as f:\n        f.write(f\"# AUTOGEN from {source_file_name} via ksc.utils.translate_and_import\")\n        f.write(py_out)\n\n    print(f.name)\n    return import_module_from_path(PYTHON_MODULE_NAME, f.name)\n\ndef subprocess_run(cmd, env=None):\n    return subprocess.run(cmd, stdout=subprocess.PIPE, env=env).stdout.decode().strip(\"\\n\")\n\ndef get_ksc_paths():\n    if \"KSC_RUNTIME_DIR\" in os.environ:\n        ksc_runtime_dir = os.environ[\"KSC_RUNTIME_DIR\"]\n    else:\n        ksc_runtime_dir = \"./src/runtime\"\n\n    if \"KSC_PATH\" in os.environ:\n        ksc_path = os.environ[\"KSC_PATH\"]\n    else:\n        ksc_path = \"./build/bin/ksc\"\n    \n    return ksc_path,ksc_runtime_dir\n\n\ndef generate_cpp_from_ks(ks_str, generate_derivatives = False, use_aten = False):\n    ksc_path,_ksc_runtime_dir = get_ksc_paths()\n\n    with NamedTemporaryFile(mode=\"w\", suffix=\".ks\", delete=False) as fks:\n        fks.write(ks_str)\n    try:\n        with NamedTemporaryFile(mode=\"w\", suffix=\".kso\", delete=False) as fkso:\n            with NamedTemporaryFile(mode=\"w\", suffix=\".cpp\", delete=False) as fcpp:\n                print(\"generate_cpp_from_ks:\", ksc_path, fks.name)\n                e = subprocess.run([\n                    ksc_path,\n                    \"--generate-cpp\" if generate_derivatives else \"--generate-cpp-without-diffs\",\n                    \"--ks-source-file\", \"src/runtime/prelude.ks\",\n                    *((\"--ks-source-file\", \"src/runtime/prelude-aten.ks\") if use_aten else ()),\n                    \"--ks-source-file\", fks.name,\n                    \"--ks-output-file\", fkso.name,\n                    \"--cpp-output-file\", fcpp.name\n                ], capture_output=True, check=True)\n                print(e.stdout.decode('ascii'))\n                print(e.stderr.decode('ascii'))\n    except subprocess.CalledProcessError as e:\n        print(f\"files {fks.name} {fkso.name} {fcpp.name}\")\n        print(f\"ks_str=\\n{ks_str}\")\n        print(e.output.decode('ascii'))\n        print(e.stderr.decode('ascii'))\n        raise\n\n\n    # Read from CPP back to string\n    with open(fcpp.name) as f:\n        out = f.read()\n\n    # only delete these file if no error\n    os.unlink(fks.name)\n    os.unlink(fcpp.name)\n    os.unlink(fkso.name)\n\n    return out\n\ndef build_py_module_from_cpp(cpp_str, pybind11_path, use_aten=False):\n    _ksc_path,ksc_runtime_dir = get_ksc_paths()\n\n    with NamedTemporaryFile(mode=\"w\", suffix=\".cpp\", delete=False) as fcpp:\n        fcpp.write(cpp_str)\n\n    extension_suffix = sysconfig.get_config_var('EXT_SUFFIX')\n    if extension_suffix is None:\n        extension_suffix = sysconfig.get_config_var('SO')\n\n    with NamedTemporaryFile(mode=\"w\", suffix=extension_suffix, delete=False) as fpymod:\n        pass\n    module_path = fpymod.name\n    module_name = os.path.basename(module_path).split(\".\")[0]\n    python_includes = subprocess_run(\n        [sys.executable, \"-m\", \"pybind11\", \"--includes\"],\n        env={\"PYTHONPATH\": \"extern/pybind11\"}\n    )\n    try:\n        cmd = (f\"g++ -I{ksc_runtime_dir} -I{pybind11_path}/include \"\n               + python_includes\n               + \" -Wall -Wno-unused-variable -Wno-unused-but-set-variable\"\n                 \" -fmax-errors=1\"\n                 \" -std=c++17\"\n                 \" -O3\"\n                 \" -fPIC\"\n                 \" -shared\"\n               + (\" -DKS_INCLUDE_ATEN\" if use_aten else \"\")\n               + f\" -DPYTHON_MODULE_NAME={module_name}\"\n                 f\" -o {module_path} \"\n               + fcpp.name)\n        print(cmd)\n        subprocess.run(cmd, shell=True, capture_output=True, check=True)\n    except subprocess.CalledProcessError as e:\n        print(f\"cpp_file={fcpp.name}\")\n        print(cmd)\n        print(e.output.decode('utf-8'))\n        print(e.stderr.decode('utf-8'))\n\n        raise\n    \n    os.unlink(fcpp.name)\n    return module_name, module_path\n\ndef mangleType(ty):\n    return ty.shortstr()\n\ndef mangleTypes(tys):\n    return \"\".join(mangleType(ty) for ty in tys)\n\n\ndef encode_name(s : str) -> str:\n    # TODO: this could be faster\n    return s.\\\n        replace('@',\"$a\").\\\n        replace(',',\"$_\").\\\n        replace('.',\"$o\").\\\n        replace('[',\"$6\").\\\n        replace(']',\"$9\").\\\n        replace('<',\"$d\").\\\n        replace('>',\"$b\").\\\n        replace('*',\"$x\").\\\n        replace(':',\"$8\")\n\ndef generate_and_compile_cpp_from_ks(ks_str, name_to_call, arg_types, return_type=None, generate_derivatives=False, use_aten=False, pybind11_path=\"extern/pybind11\"):\n\n    generated_cpp_source = generate_cpp_from_ks(ks_str, generate_derivatives=generate_derivatives, use_aten=use_aten)\n\n    cpp_str = f\"\"\"\n    #include \"knossos-pybind.h\"\n    {generated_cpp_source}\n\n    \"\"\"\n\n    args_str = mangleTypes(arg_types)\n    name_str = encode_name(f\"{name_to_call}@{args_str}\")\n    declarations = f\"\"\"\n     m.def(\"entry\", with_ks_allocator(&ks::{name_str}));\n    \"\"\"\n\n    if generate_derivatives:\n        darg_types = [tangent_type(ty) for ty in arg_types]\n        args_tuple_str = mangleType(make_tuple_if_many(arg_types))\n        dargs_tuple_str = mangleType(make_tuple_if_many(darg_types))\n        dreturn_type_str = mangleType(tangent_type(return_type))\n\n        fwd_name = encode_name(f\"fwd${name_to_call}@{args_tuple_str}\")\n        declarations += f\"\"\"\n          m.def(\"fwd_entry\", with_ks_allocator(&ks::{fwd_name}));\n        \"\"\"\n\n        rev_name = encode_name(f\"rev${name_to_call}@{args_tuple_str}\")\n        declarations += f\"\"\"\n          m.def(\"rev_entry\", with_ks_allocator(&ks::{rev_name}));\n        \"\"\"\n\n    cpp_str += \"\"\"\nPYBIND11_MODULE(PYTHON_MODULE_NAME, m) {\n    m.def(\"reset_allocator\", []{g_alloc.reset();});\n    m.def(\"allocator_top\", []{ return g_alloc.mark();});\n    m.def(\"allocator_peak\", []{ return g_alloc.peak();});\n\n    declare_tensor_2<double>(m, \"Tensor_2_Float\");\n    declare_tensor_2<int>(m, \"Tensor_2_Integer\");\n\n\"\"\" + declarations + \"\"\"\n}\n\"\"\"\n    pybindcpppath = \"obj/pybind.cpp\"\n    print(f\"Saving to {pybindcpppath}\")    \n    os.makedirs(os.path.dirname(pybindcpppath), exist_ok=True)\n    with open(pybindcpppath, \"w\") as fcpp:\n        fcpp.write(cpp_str)\n\n    module_name, module_path = build_py_module_from_cpp(cpp_str, pybind11_path, use_aten=use_aten)\n    return import_module_from_path(module_name, module_path)\n\ndef ndgrid_inds(sz):\n    \"\"\"\n    Return a sequnce of tuples of indices as if generated by nested comprehensions.\n    Example:\n        ndgrid_inds((ni,nj))\n    Returns the same sequence as\n        [(i,j) for i in range(ni) for j in range(nj)]\n\n    The iterates are always tuples so\n        ndgrid_inds(4)\n    returns\n        [(0,), (1,), (2,), (3,)] \n\n    \"\"\"\n\n    return itertools.product(*map(range, sz))\n\n", "path": "src/python/ksc/utils.py"}]}