float(x) does not work if x is a concrete tracer, but int(x) works

I expected the following code to work:

 def func(x):
    return x if float(x) >= 0. else -x
 print(api.grad(func)(3.))

because int(x) works. Instead that code gives an error

TypeError: JAX Tracer object cannot be interpreted as a float. Try using `x.astype(float)` instead.

Using x.astype(float) works, but I am wondering if it would not be more uniform to support float(x) as well.

[gnecula]

Thanks for explaining. I do think that this is a subtle issue, I do not know if people run into it in practice, or just when poking curiously at JAX. Do you think it is worth an entry in the FAQ?

WDYT? In particular, do you think we should make float(x) behave exactly like x.astype(‘float’), at least on Tracer instances?

I think that x.astype(float) will produce a tracer if x is a tracer, which would not make sense as the result of float(x).

[mattjj]

Hmm hard to say: what’s the bar for adding to the FAQ? We could just wait until it’s asked before adding it.

(I just remembered, based on a JAXers thread and a realization just now, that the original reason for adding the __float__ override may have been for printing with %f.)

I defer to your judgement on the FAQ, though my own bias is towards inaction 😃