<p dir="auto">I get the warning message<br>
<code class="notranslate">W external/org_tensorflow/tensorflow/compiler/xla/service/hlo_pass_fix.h:49] Unexpectedly high number of iterations in HLO passes, exiting fixed point loop.</code><br>
each time I evaluate <code class="notranslate">grad(odeint)</code> in the following <a href="https://github.com/jessebett/jax/blob/hk-nodes/examples/nodes_cifar_hk.py">code</a>. The forward system has <code class="notranslate">4 * 4 * 64 = 1024</code> dimensions, and the dynamics function has <code class="notranslate">1 * 1 * 64 * 64 = 4096</code> parameters (so the adjoint system would have about <code class="notranslate">5121</code> dimensions). I also <a href="https://github.com/jessebett/jax/blob/hk-nodes/jax/experimental/ode.py">modified</a> <code class="notranslate">odeint</code> to count NFE on the forward pass (which in this case is only <code class="notranslate">7</code>).</p>
<p dir="auto">Not really sure what this warning message means, but I suspect it might be affecting the performance of my code. Any insights into this would be greatly appreciate, and I'm happy to create a more stripped down example (e.g. without the haiku dependency) if needed.</p>
<p dir="auto"><a href="https://drive.google.com/open?id=13Zck0wpfd57ZtoaKfpLfZIMw5UITsIH4" rel="nofollow">XLA dump</a></p>
<p dir="auto">I am getting an error/warning message when using clip_grads. The message is:<br>
"external/org_tensorflow/tensorflow/compiler/xla/service/hlo_pass_fix.h:49] Unexpectedly high number of iterations in HLO passes, exiting fixed point loop."</p>
<p dir="auto">Does anyone know what this means, or how to call clip_grads so it doesn't occur? Does it affect any results?</p>
<p dir="auto">In the code, with option == 1 (no jit for clip_grads), the message occurs five times. The number is dependent on the number of dense layers in the neural network. If option==1 (with jit), the message occurs only once. I assume the others are suppressed.</p>
<p dir="auto">Incidently, apart from clip_grads, I also get the message (once) if I change the neural net to one input and send X[:,0] to it in net_apply().</p>
<p dir="auto">And I also get the same message when I call odeint().</p>
<p dir="auto">I'm using jax version '0.1.62' and jaxlib version '0.1.43' on the GPU.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="from jax import grad
from jax import random
from jax import jit 

from jax.experimental import stax 
from jax.experimental.stax import Dense, Tanh,  BatchNorm
from jax.experimental import optimizers

# setup neural network -------------------------------
L = 2
net_init, net_apply = stax.serial(
    Dense(L), Tanh,
    Dense(L), Tanh,
    Dense(2)
    )

rng = random.PRNGKey(0)
rng, subkey = random.split(rng)

out_shape, net_params = net_init(subkey, (-1, 2))  
opt_init, opt_update, get_params = optimizers.adam(step_size=.001)
opt_state = opt_init(net_params)

# option for jit clip_grads --------------------------
def clip(gr):
    return optimizers.clip_grads(gr, 4)

clip_jit = jit(clip)

# loss function --------------------------------------
def loss(params, X):
    ans = net_apply(params, X[:,0:2])
    return 1.  # return dummy loss

# step function --------------------------------------
def step(opt_state, X):
    params = get_params(opt_state)
    
    gr = grad(loss, argnums=(0,)) (params, X)
    print(len(params))
    # try different options
    option = [0,1,2][2]
    if option == 0:
        gr = optimizers.clip_grads(gr, 4)
    
    elif option == 1:
        gr = clip_jit(gr)
    
    elif option == 2:
        pass
    
    return 1.  # return dummy value

# setup data -----------------------------------------
rng, subkey = random.split(rng)
N = 50
X = random.normal(subkey, (N, 10))

# call step ------------------------------------------
for ii in range(2):
    result = step(opt_state, X)
"><pre class="notranslate"><code class="notranslate">from jax import grad
from jax import random
from jax import jit 

from jax.experimental import stax 
from jax.experimental.stax import Dense, Tanh,  BatchNorm
from jax.experimental import optimizers

# setup neural network -------------------------------
L = 2
net_init, net_apply = stax.serial(
    Dense(L), Tanh,
    Dense(L), Tanh,
    Dense(2)
    )

rng = random.PRNGKey(0)
rng, subkey = random.split(rng)

out_shape, net_params = net_init(subkey, (-1, 2))  
opt_init, opt_update, get_params = optimizers.adam(step_size=.001)
opt_state = opt_init(net_params)

# option for jit clip_grads --------------------------
def clip(gr):
    return optimizers.clip_grads(gr, 4)

clip_jit = jit(clip)

# loss function --------------------------------------
def loss(params, X):
    ans = net_apply(params, X[:,0:2])
    return 1.  # return dummy loss

# step function --------------------------------------
def step(opt_state, X):
    params = get_params(opt_state)
    
    gr = grad(loss, argnums=(0,)) (params, X)
    print(len(params))
    # try different options
    option = [0,1,2][2]
    if option == 0:
        gr = optimizers.clip_grads(gr, 4)
    
    elif option == 1:
        gr = clip_jit(gr)
    
    elif option == 2:
        pass
    
    return 1.  # return dummy value

# setup data -----------------------------------------
rng, subkey = random.split(rng)
N = 50
X = random.normal(subkey, (N, 10))

# call step ------------------------------------------
for ii in range(2):
    result = step(opt_state, X)

</code></pre></div>