Make pip-installable and simplify commands in README

README.md

@@ -17,23 +17,14 @@ This is not an officially supported Google product.
 
 ## Installation
 
-Installation is currently a bit manual. First, install dependencies:
-
-```
-pip3 install chex einops dm-haiku networkx typing_extensions
-```
-
-Second, clone the repo:
+Just clone and pip install:
 
 ```
 git clone https://github.com/deepmind/tracr
+cd tracr
+pip3 install .
 ```
 
-Third, put the resulting folder somewhere in your `PYTHONPATH`
-(eg by placing the `tracr` checkout in the root of your project folder).
-
-This will be made easier in the future.
-
 
 ## Usage example: RASP `reverse` program
 
@@ -190,19 +181,5 @@ First, install test dependencies:
 pip3 install absl-py pytest
 ```
 
-```
-# We use `python3 -m pytest` instead of just `pytest` so that the working directory is
-# added to PYTHONPATH.
-# -ra: Report names of tests that failed, were skipped, etc.
-python3 -m pytest -ra
-```
-
-This should take about 60 seconds. If you install `pytest-xdist`, you can run them in
-parallel with:
-
-```
-python3 -m pytest -ra -n auto
-```
-
-However, currently this only shaves off about 10 seconds, since it's bottlenecked by a
-single long-running test.
+Then, in the checkout directory, simply run `pytest`. This should take about 60
+seconds.

setup.py

@@ -0,0 +1,36 @@
+# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Install script."""
+
+import setuptools
+
+setuptools.setup(
+    name="tracr",
+    version="1.0.0",
+    url="https://github.com/deepmind/tracr",
+    author="DeepMind LMI team",
+    author_email="[email protected]",
+    description="Compiler from RASP to transformer weights",
+    packages=setuptools.find_packages(),
+    install_requires=[
+        "chex",
+        "einops",
+        "dm-haiku",
+        "jax",
+        "networkx",
+        "numpy",
+        "typing_extensions",
+    ],
+)

tracr/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

tracr/craft/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

tracr/craft/chamber/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

tracr/examples/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

tracr/rasp/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

tracr/transformer/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

tracr/utils/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2022 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================