Extracting Data from traced ConcreteArray

[vasilavramov]

How do I convert an array of type traced ConcreteArray to a normal numpy array? The normal np.array method fails. I.e. Is there any efficient way to extract the raw data from the concrete array?

[jakevdp]

You can use the val attribute of the ConcreteArray to extract the concrete value; for example:

import jax.numpy as jnp
from jax import grad

extracted_value = None

def add(a, b):
  global extracted_value
  extracted_value = a.val
  return jnp.add(a, b)

_ = grad(add)(2.0, 10.)

print(extracted_value)
# 2.0

However, for most uses of JAX you should not need to do this. Can you say more about the application that inspired this question?

[mattjj]

Indeed, opening up JAX data structures like this is likely to lead to broken code. For example:

In [1]: from jax import grad
In [2]: import jax.numpy as jnp
In [3]: grad(lambda x: jnp.cos(x.val))(3.)
Out[3]: array(0., dtype=float32)

[vasilavramov]

The goal was using numpy functions (mainly numpy arrays as they allow index allocation of data) inside a grad-ed function that get its input from a Neural network. I am aware that this is generally bad practice and ideally I need to rewrite methods so that they can be vectorised to work on the objects I am iterating over.

[mattjj]

You should probably use jax.lax.stop_gradient instead, which is a lot safer and more robust than opening up JAX-internal data structures:

In [1]: import numpy as np
In [2]: import jax
In [3]: jax.grad(lambda x: np.sin(jax.lax.stop_gradient(x)))(3.)
Out[3]: array(0., dtype=float32)

Using jax.lax.stop_gradient will peel off any autodiff Tracers and just give you a raw array.