How to use xla_bridge_get_backend().live_arrays()

[ToshiyukiBandai]

Hi JAX community,

Inspired by #18330, I tested xla_bridge_get_backend().live_arrays(), but I am quite confused by the output. If I run the following codes on a Jupyter notebook (with GPU; jax version 0.4.19)

import jax
from jax import jit, vmap, lax, jacfwd, jacrev, grad, vjp, jvp, random
import jax.numpy as jnp
from jax.lib import xla_bridge
from jax.config import config
config.update("jax_enable_x64", True)

alpha = jnp.array(4.0)
beta = jnp.array(1.0)
gamma = jnp.array(4.0)
delta = jnp.array(4.0)

x_test = jnp.array([1.0, 3.0])
y_test = jnp.array([2.0, 3.0])

print(xla_bridge.get_backend().live_arrays())

I got the output

[Array([2., 3.], dtype=float64), Array([1., 3.], dtype=float64), Array(4., dtype=float64, weak_type=True), Array(4., dtype=float64, weak_type=True), Array(1., dtype=float64, weak_type=True), Array(4., dtype=float64, weak_type=True)]

which is good. But if I print the output again, I got

[Array(4., dtype=float64, weak_type=True), Array(1., dtype=float64, weak_type=True), Array(4., dtype=float64, weak_type=True), Array(4., dtype=float64, weak_type=True), Array([1., 3.], dtype=float64), Array([2., 3.], dtype=float64), Array([2., 3.], dtype=float64), Array([1., 3.], dtype=float64), Array(4., dtype=float64, weak_type=True), Array(4., dtype=float64, weak_type=True), Array(1., dtype=float64, weak_type=True), Array(4., dtype=float64, weak_type=True)]

The arrays were doubled! Does anyone know what is going on here?

[hawkinsp]

Please use jax.live_arrays(), which is the public API.

When you say "print the output again": did you run the entire Jupyter cell again? Is it possible that Jupyter is keeping references to the first set of outputs?

[ToshiyukiBandai]

Understood, I switched to jax.live_arrays() but got the same results.

I did not run the entire Jupyter cell again, I just ran an additional print statement. I modified the code and confirmed that JAX somehow creates new copies.

import jax
from jax import jit, vmap, lax, jacfwd, jacrev, grad, vjp, jvp, random
import jax.numpy as jnp
from jax.lib import xla_bridge
from jax.config import config
config.update("jax_enable_x64", True)

alpha = jnp.array(1.0)
beta = jnp.array(2.0)
gamma = jnp.array(3.0)
delta = jnp.array(4.0)

x_test = jnp.array([1.0, 3.0])
y_test = jnp.array([2.0, 3.0])

jax.__version__

x = jax.live_arrays()
x2 = jax.live_arrays()

x[-1] # refers to alpha
x2[0] # refers to alpha?
x2[-1] # refers to alpha?

id(x[-1]) == id(x2[0]) # False!
id(x[-1]) == id(x2[-1]) # True!

[jakevdp]

I'm not sure how you're executing your code, but note that Jupyter notebooks store references to all cell outputs. So the fact that you wrote things like x[-1] # refers to alpha makes me think that you're displaying this in a cell output, which would maintain a reference to it. You can use print(x[-1]) to avoid this.

[ToshiyukiBandai]

Thank you, you're right. I confirmed jax.live_arrays() showed the right things (no copies).