Repository: google/jax

Issue:
Fix large constants in activations
This fixes an issue where using the elu, selu, or tanh activation results in a constant as big as the input being created.

This PR also switches the use of np.where to lax.select for activation functions to avoid accidental broadcasting of scalars in the future.

I added a test that verifies that making the jaxpr does not produce constants that scale in size with the input.


Discussion:
Thanks for this! Looks like there were some test failures in 64bit mode.
Turns out lax.select fails because in X64 mode a python float * float32 tensor = float64 tensor. Is that not a violation of numpy dtype promotion rules?
I'm hesitant to merge this just because it makes the code uglier and we hope #1668 (or a related PR) will fix the underlying issue. But at the same time this is fixing a real problem!

@jekbradbury WDYT? Should we merge, maybe with TODO notes to revert it once we fix the underlying issues?
I think I’d rather merge a version that uses `tie_in` to avoid materializing broadcast constants, if we can get that to work (I can look into it), or wait until #1668 if we expect that to land soon.

The version in this PR depends on the specific current behavior where lax.full is lazy but broadcasts aren’t, which makes it a little opaque.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -1337,8 +1337,8 @@ def jaxpr_maker(*args, **kwargs):
     jax_args, in_tree = tree_flatten((args, kwargs))
     jaxtree_fun, out_tree = flatten_fun(wrapped, in_tree)
     in_pvals = map(pv_like, jax_args)
-    jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jaxtree_fun, in_pvals,
-                                                 instantiate=True)
+    jaxpr, out_pvals, consts = pe.trace_to_jaxpr(
+        jaxtree_fun, in_pvals, instantiate=True, stage_out_calls=True)
     out_avals = map(raise_to_shaped, unzip2(out_pvals)[0])
     in_avals = tuple(raise_to_shaped(in_aval) for in_aval, _ in in_pvals)
     typed_jaxpr = core.TypedJaxpr(jaxpr, consts, in_avals, out_avals)
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -325,10 +325,10 @@ def partial_val_aval(pv, const):
   else:
     raise TypeError(pv)
 
-def trace_to_jaxpr(fun, pvals, **kwargs):
+def trace_to_jaxpr(fun, pvals, instantiate=False, stage_out_calls=False):
   """Traces a function, given abstract inputs, to a jaxpr."""
-  instantiate = kwargs.pop('instantiate', False)
-  with new_master(JaxprTrace) as master:
+  trace_type = StagingJaxprTrace if stage_out_calls else JaxprTrace
+  with new_master(trace_type) as master:
     fun = trace_to_subjaxpr(fun, master, instantiate)
     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
     assert not env
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -57,7 +57,8 @@
 def _initial_style_jaxpr(fun, in_tree, in_avals):
   in_pvals = [pe.PartialVal((aval, core.unit)) for aval in in_avals]
   fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
-  jaxpr, out_pvals, consts = pe.trace_to_jaxpr(fun, in_pvals, instantiate=True)
+  jaxpr, out_pvals, consts = pe.trace_to_jaxpr(fun, in_pvals, instantiate=True,
+                                               stage_out_calls=True)
   out_avals = _map(raise_to_shaped, unzip2(out_pvals)[0])
   const_avals = tuple(raise_to_shaped(core.get_aval(c)) for c in consts)
   typed_jaxpr = core.TypedJaxpr(pe.closure_convert_jaxpr(jaxpr),


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -953,6 +953,14 @@ def test_issue_871(self):
   def test_partial_eval_lower(self):
     # this is a simplified model of a bug that arose when we first used @jit in
     # a jvp rule. it's in this file because we want to use make_jaxpr.
+
+    # NOTE(mattjj): I no longer understand what this was meant to test. My guess
+    # is it was related to staging out the broadcast into a jaxpr to be
+    # transposed, but after #1749 that's no longer a problem. After changing
+    # make_jaxpr (and jit) to stage out sub-calls fully, this test started to
+    # fail; I left it in as skipped because deleting tests feels wrong.
+    raise unittest.SkipTest("obsolete test")
+
     @api.jit
     def f(a, b, c):
       a = lax.broadcast(a, (2,))
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -30,6 +30,7 @@
 from jax import nn
 from jax import random
 import jax
+import jax.numpy as np
 
 from jax.config import config
 config.parse_flags_with_absl()
@@ -51,6 +52,16 @@ def testEluValue(self):
     val = nn.elu(1e4)
     self.assertAllClose(val, 1e4, check_dtypes=False)
 
+  @jtu.skip_on_devices("gpu", "tpu")
+  def testEluMemory(self):
+    # see https://github.com/google/jax/pull/1640
+    jax.make_jaxpr(nn.elu)(np.ones((10 ** 12,)))  # don't oom
+
+  @jtu.skip_on_devices("gpu", "tpu")
+  def testHardTanhMemory(self):
+    # see https://github.com/google/jax/pull/1640
+    jax.make_jaxpr(nn.hard_tanh)(np.ones((10 ** 12,)))  # don't oom
+
 InitializerRecord = collections.namedtuple(
   "InitializerRecord",
   ["name", "initializer", "shapes"])