jax.Array.view returns a copy rather than a view

[hirayaku]

I would like to reinterpret a jax array of dtype uint32 to uint8. While .view(dtype=jnp.uint8) does this job, the doc claims it performs a bitwise copy first:
https://jax.readthedocs.io/en/latest/_autosummary/jax.Array.html#jax.Array.__init__:~:text=Return%20a-,bitwise%20copy,-of%20the%20array

This part of the code is performance critical and I want to avoid copies as much as I could. Is there another API that actually returns a reinterpreted view?

[jakevdp]

Hi - thanks for the question! This is working as intended. Two JAX arrays can never share the same memory, and so there is no way for view to produce a second array that shares the same memory as the first.

Keep in mind, however, that within a jit context this restriction is lifted, and the compiler can (and generally will) choose a more memory-efficient approach when it is possible to do so without violating the semantics of the original program.

[hirayaku]

I was wondering which operation counts as violating the semantics. Is it documented anywhere and is read-only reinterpret okay? Thanks.

[jakevdp]

I was wondering which operation counts as violating the semantics.

for example, a program like this:

@jit
def f(x):
  return x.view('uint8')

Because two materialized JAX arrays may not share the same memory, this function will always create a new buffer for y.

On the other hand, something like this:

@jit
def f():
  x = jnp.arange(10)
  return x.view('uint8')

it would not violate JAX's semantics to create the view in an inplace fashion, because only one materialized array exists.