Bitcast array from int32 to int64

[juliandwain]

Hi,
I'm trying to replicate the behavior of tensorflow's bitcast function.
I have an array of type int32 and I want to convert it to int64.
I found that jax has the jax.lax.bitcast_convert_type function, but this does not work properly.
When I am wrapping my operation inside jax.experimental.enable_x64, I get the following error:

Exception has occurred: TypeError
`bitcast_convert_type` for operand type (int32) must have destination type (int64) of same size.
  File "/home/user/proj/gendata.py", line 47, in function
    e= jax.lax.bitcast_convert_type(
  File "/home/user/proj/gendata.py", line 60, in <module>
    a, b = function1(inputs)

For context:

• I have an input array a with shape (100,2) with dtype int32
• according to the tensorflow documentation, tf.bitcast, "If T is smaller than type, the operator requires that the rightmost dimension be equal to sizeof(type)/sizeof(T). The shape then goes from [..., sizeof(type)/sizeof(T)] to [...].", in my case, the array is flattened
• the wrapping is done like such:

with jax.experimental.enable_x64():
  a = jax.lax.bitcast_convert_type(...)

As I am currently pretty much stuck here, I appreciate every little help I get!
Essentially, I am trying to replicate tf.bitcast behavior in jax.

Thanks in advance!

Best,
J

[jakevdp]

One way to do this is with the view method of JAX arrays. For example:

import jax
jax.config.update('jax_enable_x64', True)

import jax.numpy as jnp
x32 = jnp.arange(20, dtype='int32').reshape(10, 2)
x64 = x32.view('int64').ravel()
print(x64)
# [ 4294967296 12884901890 21474836484 30064771078 38654705672 47244640266
#  55834574860 64424509454 73014444048 81604378642]

import tensorflow as tf
x32 = tf.reshape(tf.range(20, dtype='int32'), (10, 2))
x64 = tf.bitcast(x32, 'int64')
print(x64)
# tf.Tensor(
# [ 4294967296 12884901890 21474836484 30064771078 38654705672 47244640266
#  55834574860 64424509454 73014444048 81604378642], shape=(10,), dtype=int64)

[juliandwain]

That will give you a (*, 1) array. You could get the original back with x64[:,, None].view('int32')

Yes sorry that was a typo from my side, I meant 2 not 1

[juliandwain]

I think the issue is within the unique call, since the return values are different for JAX and TF:

JAX

[ 4294967296 12884901890 21474836484 30064771078 38654705672 47244640266
 55834574860 64424509454 73014444048 81604378642]

TF

[38654705672 73014444048 12884901890 30064771078 64424509454 81604378642
 55834574860  4294967296 21474836484 47244640266]

They contain the same values but the ordering is different.

When executing

_arr = jnp.asarray(
    [
        38654705672,
        73014444048,
        12884901890,
        30064771078,
        64424509454,
        81604378642,
        55834574860,
        4294967296,
        21474836484,
        47244640266,
    ]
)
print(_arr[:, None].view("int32"))
# [[ 8  9]
 # [16 17]
 # [ 2  3]
 # [ 6  7]
 # [14 15]
 # [18 19]
 # [12 13]
 # [ 0  1]
 # [ 4  5]
 #[10 11]]

the output is exactly the same as the one from tensorflow.

[jakevdp]

RIght, jnp.unique is designed to have the same output as np.unique, which sorts its outputs.

[juliandwain]

Is it possible to make it similar to TF?

[jakevdp]

I think this should give the same result as tf.unique, though I haven't tested it extensively:

x = jnp.array([4, 2, 4, 3, 3, 2, 4])
x_unique, i = jnp.unique(x, return_index=True)
print(x_unique[jnp.argsort(i)])
# [4 2 3]