<h3 dir="auto">Description</h3>
<p dir="auto">Seeing an issue where trying to jit() run the function results in the following error:</p>
<p dir="auto"><code class="notranslate">XlaRuntimeError: INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/hlo/ir/hlo_computation.cc:962) ShapeUtil::Compatible(old_instruction-&gt;shape(), new_instruction-&gt;shape()) f32[1,5,3,12] vs f32[1,5,7,12]</code></p>
<p dir="auto">I've narrowed it down to a combination of:</p>
<ul dir="auto">
<li>Two Conv2D:s on the same underlying array</li>
<li>Recombining them in some fashion (e.g. addition) and blocking gradient flow through one of these packs with stop_gradient</li>
<li>Doing a linear projection on the result</li>
<li>Taking grad() of all of the above.</li>
<li>Running on GPU (succeeds on CPU).</li>
</ul>
<p dir="auto">See notebook where the non-jit version succeeds, but jit() of the same function fails to execute with an internal error.</p>
<p dir="auto"><a href="https://colab.research.google.com/drive/1QybGmVwN90UIg7TXgqmwsS9QSiz-GrH7" rel="nofollow">https://colab.research.google.com/drive/1QybGmVwN90UIg7TXgqmwsS9QSiz-GrH7</a></p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">0.3.25</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">GPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Colab GPU kernel</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto">Thu Dec 29 15:45:31 2022<br>
+-----------------------------------------------------------------------------+<br>
| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |<br>
|-------------------------------+----------------------+----------------------+<br>
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |<br>
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |<br>
|                               |                      |               MIG M. |<br>
|===============================+======================+======================|<br>
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |<br>
| N/A   44C    P0    27W /  70W |  13610MiB / 15109MiB |      0%      Default |<br>
|                               |                      |                  N/A |<br>
+-------------------------------+----------------------+----------------------+</p>
<p dir="auto">+-----------------------------------------------------------------------------+<br>
| Processes:                                                                  |<br>
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |<br>
|        ID   ID                                                   Usage      |<br>
|=============================================================================|</p>
<p dir="auto">Hello,</p>
<p dir="auto">I am trying to understand exactly how auto-diff works with JAX. I have an example function:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def fnc_jax(x1, x2):
    
    return (jnp.divide(x1,x2) - jnp.exp(x2))*(jnp.sin(jnp.divide(x1,x2)) + jnp.divide(x1,x2) - jnp.exp(x2))"><pre class="notranslate"><code class="notranslate">def fnc_jax(x1, x2):
    
    return (jnp.divide(x1,x2) - jnp.exp(x2))*(jnp.sin(jnp.divide(x1,x2)) + jnp.divide(x1,x2) - jnp.exp(x2))
</code></pre></div>
<p dir="auto">Now I do <code class="notranslate">jax.make_jaxpr(fnc_jax)(1.0,1.0)</code>, which gives me the following output:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="{ lambda ; a:f32[] b:f32[]. let
    c:f32[] = div a b
    d:f32[] = exp b
    e:f32[] = sub c d
    f:f32[] = div a b
    g:f32[] = sin f
    h:f32[] = div a b
    i:f32[] = add g h
    j:f32[] = exp b
    k:f32[] = sub i j
    l:f32[] = mul e k
  in (l,) }"><pre class="notranslate"><code class="notranslate">{ lambda ; a:f32[] b:f32[]. let
    c:f32[] = div a b
    d:f32[] = exp b
    e:f32[] = sub c d
    f:f32[] = div a b
    g:f32[] = sin f
    h:f32[] = div a b
    i:f32[] = add g h
    j:f32[] = exp b
    k:f32[] = sub i j
    l:f32[] = mul e k
  in (l,) }
</code></pre></div>
<p dir="auto">Question 1: Why is it that c, e, and h for example are doing the exact same computation? Wouldn't a better way be if we had something like this:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="def fnc_jax_alt(x1, x2):
    a = x1/x2      
    b = np.exp(x2)     
    c = np.sin(a)
    d = a - b
    e = c + d
    g = d * e
    return g"><pre class="notranslate"><code class="notranslate">def fnc_jax_alt(x1, x2):
    a = x1/x2      
    b = np.exp(x2)     
    c = np.sin(a)
    d = a - b
    e = c + d
    g = d * e
    return g
</code></pre></div>
<p dir="auto">which gives the same output for a given input as <code class="notranslate">fnc_jax</code>, but doesnt repeat computations confirmed using <code class="notranslate">jax.make_jaxpr(fnc_jax_alt)(1.0,1.0)</code>. Is this a feature or a bug?</p>
<p dir="auto">Another confusion I have is when I do <code class="notranslate">jax.make_jaxpr(jax.grad(fnc_jax))(1.0,1.0)</code>:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="{ lambda ; a:f32[] b:f32[]. let
    c:f32[] = div a b
    d:f32[] = exp b
    e:f32[] = sub c d
    f:f32[] = div a b
    g:f32[] = sin f
    h:f32[] = cos f
    i:f32[] = div a b
    j:f32[] = add g i
    k:f32[] = exp b
    l:f32[] = sub j k
    _:f32[] = mul e l
    m:f32[] = mul e 1.0
    n:f32[] = mul 1.0 l
    o:f32[] = div m b
    p:f32[] = mul m h
    q:f32[] = div p b
    r:f32[] = add_any o q
    s:f32[] = div n b
    t:f32[] = add_any r s
  in (t,) }"><pre class="notranslate"><code class="notranslate">{ lambda ; a:f32[] b:f32[]. let
    c:f32[] = div a b
    d:f32[] = exp b
    e:f32[] = sub c d
    f:f32[] = div a b
    g:f32[] = sin f
    h:f32[] = cos f
    i:f32[] = div a b
    j:f32[] = add g i
    k:f32[] = exp b
    l:f32[] = sub j k
    _:f32[] = mul e l
    m:f32[] = mul e 1.0
    n:f32[] = mul 1.0 l
    o:f32[] = div m b
    p:f32[] = mul m h
    q:f32[] = div p b
    r:f32[] = add_any o q
    s:f32[] = div n b
    t:f32[] = add_any r s
  in (t,) }
</code></pre></div>
<p dir="auto">Question 2: Apart from repeated computations, what confuses me is what exactly is going on here? It seems to me like it is doing something of a forward differentiation. But if I understand correctly from the documentation, <code class="notranslate">jax.grad</code> performs reverse-mode differentiation by default. Am I wrong or mis-understanding the make_jaxpr?</p>
<p dir="auto">If it is forward-mode differentiation, then should we not have something like the following:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="{ lambda ; a:f32[] b:f32[]. let
    c:f32[] = div a b
    d:f32[] = exp b
    e:f32[] = sin c
    f:f32[] = sub c d
    g:f32[] = add e f
    _:f32[] = mul f g
    h:f32[] = div 1.0 b
    i:f32[] = cos c
    j:f32[] = mul i h
    k:f32[] = sub h 0.0
    l:f32[] = add j k
    m:f32[] = mul k g
    n:f32[] = mul f l
    o:f32[] = add m n
  in (o,) }"><pre class="notranslate"><code class="notranslate">{ lambda ; a:f32[] b:f32[]. let
    c:f32[] = div a b
    d:f32[] = exp b
    e:f32[] = sin c
    f:f32[] = sub c d
    g:f32[] = add e f
    _:f32[] = mul f g
    h:f32[] = div 1.0 b
    i:f32[] = cos c
    j:f32[] = mul i h
    k:f32[] = sub h 0.0
    l:f32[] = add j k
    m:f32[] = mul k g
    n:f32[] = mul f l
    o:f32[] = add m n
  in (o,) }
</code></pre></div>
<p dir="auto">If it is reverse-mode differentiation, how can I get a <code class="notranslate">make_jaxpr</code> for forward-mode?</p>
<p dir="auto">Apologies in advance for my mis-understandings. I would appreciate some clarity on these questions!</p>
<p dir="auto">Regards,<br>
Jay Sandesara</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto"><em>No response</em></p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto"><em>No response</em></p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto"><em>No response</em></p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto"><em>No response</em></p>