def testReluGradientConstants(self):

    def body(i, xy):
        (x, y) = xy
        y = (y + jax.grad((lambda z: jnp.sum(jnp.maximum(z, 0.0))))(x))
        return (x, y)
    f = (lambda y: lax.fori_loop(0, 5, body, (y, y)))
    wrapped = linear_util.wrap_init(f)
    pv = partial_eval.PartialVal.unknown(jax.ShapedArray((3, 4), onp.float32))
    (_, _, consts) = partial_eval.trace_to_jaxpr(wrapped, [pv])
    self.assertFalse(any((onp.array_equal(x, onp.full((3, 4), 2.0, dtype=onp.float32)) for x in consts)))
def testReluGradientConstants(self):

    def body(i, xy):
        (x, y) = xy
        y = (y + jax.grad((lambda z: jnp.sum(jnp.maximum(z, 0.0))))(x))
        return (x, y)
    f = (lambda y: lax.fori_loop(0, 5, body, (y, y)))
    wrapped = linear_util.wrap_init(f)
    pv = partial_eval.PartialVal.unknown(jax.ShapedArray((3, 4), onp.float32))
    (_, _, consts) = partial_eval.trace_to_jaxpr(wrapped, [pv])
    self.assertFalse(any((onp.array_equal(x, onp.full((3, 4), 2.0, dtype=onp.float32)) for x in consts)))