Compare the equality of address between two arrays.

[sh0416]

x = jnp.zeros((4, 4))
y = x
print(some_equality_f(x, y))  # True because they share the same space
x = jnp.zeros((4, 4))
y = jnp.zeros((4, 4))
print(some_equality_f(x, y))  # False because x and y use different space although their values are the same.

I am certain that this functionality would be implemented in jax, but I couldn't find it.. Could you help me?

[jakevdp]

You can do this:

x = jnp.zeros((4, 4))
y = x

print(x is y)  # Python-level object equality
# True

print(x.unsafe_buffer_pointer() == y.unsafe_buffer_pointer())  # memory pointer equality
# True

But note that, as the method states, this is not safe: it is implementation-dependent, and the exact behavior may change. Also, keep in mind that this only exists in eager mode: there's no way to do this kind of check within transformations like jit, because in general the compiler needs to be free to make choices about the storage of arrays (and whether to materialize them at all).

[sh0416]

I have one more question. Following to your last sentence, compiler could copy the array that is pointed by two variables (x and y) and manage them separately to optimize their behavior whenever they needed and it doesn't break the outputs (or.. vise versa), is it right? I think it is a reasonable choice for compiler to do this automatically and it is really interesting to me that they could do this..

The reason why I want to check this equality is to guide the compiler to hint some memory efficient behavior through parameter sharing, but it maybe seems I don't have to do this..

[jakevdp]

I can't think of any way that one could use unsafe_buffer_pointer to influence the compiler. Maybe what you have in mind is the donate_argnums argument for JIT? This lets you tell the compiler which of the input buffers can be reused in the course of the computation.