Merge pull request jax-ml#2999 from gnecula/jax_outfeed_undo

Undo the id_print/id_tap feature (PR jax-ml#2791)
BuddenD · May 7, 2020 · d679ccd · d679ccd
2 parents 0a7974e + 769d703
commit d679ccd
Show file tree

Hide file tree

Showing 11 changed files with 12 additions and 1,963 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -26,7 +26,6 @@ before_install:
   - conda update -q conda
 install:
   - conda install --yes python=$TRAVIS_PYTHON_VERSION pip absl-py opt_einsum numpy scipy pytest-xdist pytest-benchmark mypy=0.770
-  - pip install msgpack
   - if [ "$JAX_ONLY_CHECK_TYPES" = true ]; then
       pip install pytype ;
     fi

diff --git a/docs/jax.experimental.host_callback.rst b/docs/jax.experimental.host_callback.rst
diff --git a/docs/jax.experimental.rst b/docs/jax.experimental.rst
@@ -4,7 +4,6 @@ jax.experimental package
 .. toctree::
     :maxdepth: 1
 
-    jax.experimental.host_callback
     jax.experimental.loops
     jax.experimental.optimizers
     jax.experimental.optix

diff --git a/jax/BUILD b/jax/BUILD
@@ -100,10 +100,3 @@ pytype_library(
     srcs_version = "PY3",
     deps = [":jax"],
 )
-
-pytype_library(
-    name = "experimental_host_callback",
-    srcs = ["experimental/host_callback.py"],
-    srcs_version = "PY3",
-    deps = [":jax"],
-)
diff --git a/jax/api.py b/jax/api.py
@@ -325,7 +325,6 @@ def computation_maker(*args, **kwargs):
     jaxpr, _, consts = pe.trace_to_jaxpr(jaxtree_fun, pvals,
                                          instantiate=instantiate_const_outputs,
                                          stage_out=True)
-    jaxpr, _ = xla.apply_outfeed_rewriter(jaxpr)
     axis_env_ = make_axis_env(xla.jaxpr_replicas(jaxpr))
     c = xb.make_computation_builder('xla_computation_{}'.format(fun_name))
     xla_consts = map(partial(xb.constant, c), consts)

diff --git a/jax/core.py b/jax/core.py
@@ -932,12 +932,7 @@ def strip_weak_type(self):
   _oct     = partialmethod(_forward_to_value, oct)
 
 
-class AbstractToken(AbstractValue):
-  def join(self, other):
-    if isinstance(other, AbstractToken):
-      return self
-    else:
-      assert False, f"Cannot join {self} with {other}"
+class AbstractToken(AbstractValue): pass
 
 abstract_token = AbstractToken()