Repository: google/jax

Issue:
AssertionError: If you see this error, please let us know by opening an issue at...
Hi there!

Just ran into 
```
AssertionError: If you see this error, please let us know by opening an issue at
https://github.com/google/jax/issues 
since we thought this was unreachable!
```
so I thought I better do as i was told!

I am trying to implement an ML algorithm, and when jit tracing a function I get the error above.

The issue was that i used a good old numpy randint inside the loss function, and when replacing that with a jax.random randint the error went away.

AssertionError: If you see this error, please let us know by opening an issue at...
Hi there!

Just ran into 
```
AssertionError: If you see this error, please let us know by opening an issue at
https://github.com/google/jax/issues 
since we thought this was unreachable!
```
so I thought I better do as i was told!

I am trying to implement an ML algorithm, and when jit tracing a function I get the error above.

The issue was that i used a good old numpy randint inside the loss function, and when replacing that with a jax.random randint the error went away.



Discussion:
Thanks for raising this! However, it's not very actionable for us without a bit more information. Can you share a repro (smaller is better)?
Yes! A branch where the error is left is here: https://github.com/jotsif/fm_jax/tree/jit_error, and run `warp_jit_error.py` to reproduce. After playing around a bit I think the error is related to that `neg_item` scalar being produced by a vanilla python function, `npr.randint`, so it catches some corner case for the 0-dimensional ConcreteArray in the jit-logic but I am on very thin ice here.

And thanks for an awesome framework btw. Got the basics of the factorisation machine working yesterday 🎉 
Thanks for the kind words, and the repro! I'll look into it.

By the way, as a workaround if you want to dodge this error, you can put this at the top of your file:

```python
from jax.interpreters import partial_eval as pe
pe._thread_local_state.remat = True
```

[The error](https://github.com/google/jax/blob/82dbf9131105a0b3a22c191930bc42b15d420794/jax/interpreters/partial_eval.py#L598-L608) is a "voluntary" and [temporary](https://github.com/google/jax/blob/82dbf9131105a0b3a22c191930bc42b15d420794/jax/interpreters/partial_eval.py#L512) one that I inserted to check our understanding that a change we made in #1749 wouldn't cause any extra (but otherwise harmless) FLOPs to be performed. The above lines just disable the "voluntary" error-checking mechanism.

I want to look into this example to understand if it's an edge case (which we can just patch) or if I should revert the extra-FLOP-potential part of #1749 that this error is meant to check.
Thanks for raising this! However, it's not very actionable for us without a bit more information. Can you share a repro (smaller is better)?
Yes! A branch where the error is left is here: https://github.com/jotsif/fm_jax/tree/jit_error, and run `warp_jit_error.py` to reproduce. After playing around a bit I think the error is related to that `neg_item` scalar being produced by a vanilla python function, `npr.randint`, so it catches some corner case for the 0-dimensional ConcreteArray in the jit-logic but I am on very thin ice here.

And thanks for an awesome framework btw. Got the basics of the factorisation machine working yesterday 🎉 
Thanks for the kind words, and the repro! I'll look into it.

By the way, as a workaround if you want to dodge this error, you can put this at the top of your file:

```python
from jax.interpreters import partial_eval as pe
pe._thread_local_state.remat = True
```

[The error](https://github.com/google/jax/blob/82dbf9131105a0b3a22c191930bc42b15d420794/jax/interpreters/partial_eval.py#L598-L608) is a "voluntary" and [temporary](https://github.com/google/jax/blob/82dbf9131105a0b3a22c191930bc42b15d420794/jax/interpreters/partial_eval.py#L512) one that I inserted to check our understanding that a change we made in #1749 wouldn't cause any extra (but otherwise harmless) FLOPs to be performed. The above lines just disable the "voluntary" error-checking mechanism.

I want to look into this example to understand if it's an edge case (which we can just patch) or if I should revert the extra-FLOP-potential part of #1749 that this error is meant to check.

Resolving patch:
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
@@ -211,7 +211,7 @@ def while_loop(cond_fun, body_fun, init_val):
   return tree_unflatten(body_tree, outs)
 
 def _while_loop_abstract_eval(*args, **kwargs):
-  return kwargs["body_jaxpr"].out_avals
+  return _map(raise_to_shaped, kwargs["body_jaxpr"].out_avals)
 
 def _while_loop_translation_rule(c, axis_env, *args, **kwargs):
   backend = kwargs.pop('backend')
@@ -363,7 +363,7 @@ def cond(pred, true_operand, true_fun, false_operand, false_fun):
   return tree_unflatten(true_out_tree, out)
 
 def _cond_abstract_eval(*args, **kwargs):
-  return kwargs["true_jaxpr"].out_avals
+  return _map(raise_to_shaped, kwargs["true_jaxpr"].out_avals)
 
 def _cond_translation_rule(c, axis_env, pred, *args, **kwargs):
   backend = kwargs.pop("backend", None)


Test patch: