Issue 1:
Description

Seeing an issue where trying to jit() run the function results in the following error:

XlaRuntimeError: INTERNAL: RET_CHECK failure (external/org_tensorflow/tensorflow/compiler/xla/hlo/ir/hlo_computation.cc:962) ShapeUtil::Compatible(old_instruction->shape(), new_instruction->shape()) f32[1,5,3,12] vs f32[1,5,7,12]

I've narrowed it down to a combination of:

Two Conv2D:s on the same underlying array

Recombining them in some fashion (e.g. addition) and blocking gradient flow through one of these packs with stop_gradient

Doing a linear projection on the result

Taking grad() of all of the above.

Running on GPU (succeeds on CPU).

See notebook where the non-jit version succeeds, but jit() of the same function fails to execute with an internal error.

https://colab.research.google.com/drive/1QybGmVwN90UIg7TXgqmwsS9QSiz-GrH7

What jax/jaxlib version are you using?

0.3.25

Which accelerator(s) are you using?

GPU

Additional system info

Colab GPU kernel

NVIDIA GPU info

Thu Dec 29 15:45:31 2022

+-----------------------------------------------------------------------------+

| NVIDIA-SMI 460.32.03 Driver Version: 460.32.03 CUDA Version: 11.2 |

|-------------------------------+----------------------+----------------------+

| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |

| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |

| | | MIG M. |

|===============================+======================+======================|

| 0 Tesla T4 Off | 00000000:00:04.0 Off | 0 |

| N/A 44C P0 27W / 70W | 13610MiB / 15109MiB | 0% Default |

| | | N/A |

+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+

| Processes: |

| GPU GI CI PID Type Process name GPU Memory |

| ID ID Usage |

|=============================================================================|

Issue 2:
Hello,

I am trying to understand exactly how auto-diff works with JAX. I have an example function:

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))

Now I do jax.make_jaxpr(fnc_jax)(1.0,1.0), which gives me the following output:

{ 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,) }

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:

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

which gives the same output for a given input as fnc_jax, but doesnt repeat computations confirmed using jax.make_jaxpr(fnc_jax_alt)(1.0,1.0). Is this a feature or a bug?

Another confusion I have is when I do jax.make_jaxpr(jax.grad(fnc_jax))(1.0,1.0):

{ 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,) }

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, jax.grad performs reverse-mode differentiation by default. Am I wrong or mis-understanding the make_jaxpr?

If it is forward-mode differentiation, then should we not have something like the following:

{ 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,) }

If it is reverse-mode differentiation, how can I get a make_jaxpr for forward-mode?

Apologies in advance for my mis-understandings. I would appreciate some clarity on these questions!

Regards,

Jay Sandesara

What jax/jaxlib version are you using?

No response

Which accelerator(s) are you using?

No response

Additional system info

No response

NVIDIA GPU info

No response