convert tracer to numpy array?

[dui1234]

Hi,

I am struggling with this conversion of tracer to numpy array. I have a function that calculate a loss:

@jit
def smooth_l1_loss_none(x,y,beta=1.0):
    ln1 = lambda x,y:0.5*(x-y)**2/beta
    ln2 = lambda x,y:jnp.abs(x-y)-0.5*beta
    def func(x,y):
        cond = jnp.abs(x-y) < beta
        return jnp.where(cond,ln1(x,y),ln2(x,y))
    L = vmap(func)(x,y)
    return L

and I want to convert the product of this function to a numpy array via:
loss_for_prior = np.array(elementwise_loss)

However, it gives me an error of

TracerArrayConversionError: The numpy.ndarray conversion method array() was called on the JAX Tracer object Traced.....

What should I do in this case?

[dui1234]

additional info:

the function is called in the following manner:

loss, grads = jax.value_and_grad(self._compute_2_losses)(params,param2,samples1,
                                                                 samples2,weights,indices)

    @partial(jit, static_argnums=(0,))
    def _compute_dqn_loss(self,params1,params2,samples,gamma):
        """Return dqn loss."""
        state = jnp.array(samples["obs"])
        next_state = jnp.array(samples["next_obs"])
        action = jnp.array(samples["acts"].reshape(-1, 1),dtype=jnp.int32)
        reward = jnp.array(samples["rews"].reshape(-1, 1))
        done = jnp.array(samples["done"].reshape(-1, 1))

        # G_t   = r + gamma * v(s_{t+1})  if state != Terminal
        #       = r                       otherwise
        gather = vmap(lambda x,i:x[i])
        curr_q_value = gather(self.dqn[1](params1,state),action)
        next_q_value = gather(self.dqn_target[1](params2,next_state),
                     self.dqn[1](params1,next_state).argmax(axis=1,keepdims=True))
        mask = 1 - done
        target = (reward + gamma * next_q_value * mask)

        # calculate dqn loss
        elementwise_loss = smooth_l1_loss_none(curr_q_value, target)

        return elementwise_loss
    
    @partial(jit, static_argnums=(0,))
    def _compute_2_losses(self,params1,params2,samples1,samples2,weights,indices):
        
        # 1-step Learning loss
        elementwise_loss = self._compute_dqn_loss(params1,params2,samples1, self.gamma)
        loss = jnp.mean(elementwise_loss * weights)
        
        # N-step Learning loss
        if self.use_n_step:
            gamma = self.gamma ** self.n_step
            elementwise_loss_n_loss = self._compute_dqn_loss(params1, params2, samples2, gamma)
            elementwise_loss += elementwise_loss_n_loss
            loss = jnp.mean(elementwise_loss * weights)
        
        # PER: update priorities
        loss_for_prior = np.array(elementwise_loss)
        new_priorities = loss_for_prior + self.prior_eps
        self.memory.update_priorities(indices, new_priorities)
        
        return loss

I suspect that it might be something to do with the jax.value_and_grad as the function is called under this?? but I dont know exactly how should I solve this.

[jakevdp]

Edit: there is now a FAQ entry on this topic: https://jax.readthedocs.io/en/latest/faq.html#how-can-i-convert-a-jax-tracer-to-a-numpy-array

---

You cannot convert a JAX tracer to a numpy array: a JAX tracer is an abstract value that represents all possible arrays of a given shape and dtype; a numpy array is a concrete value representing one particular array.

It sounds like what you want to do is to call non-JAX code from within transformed JAX code. This is possible to do via jax.pure_callback (see also External Callbacks in JAX) but before going this route you should be aware that it can have severe performance penalties, particularly if you're running on accelerators like GPU.

Can you say more about why you want to convert these values to numpy arrays? Perhaps there's a way to keep your computation in JAX rather than calling non-JAX libraries?

[dui1234]

Hi Jake,

Thanks for your reply. I need the conversion because my self.memory object is written using numpy and the method self.memory.update_priorities takes numpy arrays/lists as the argument.

One solution I can think of is to move the whole part of

# PER: update priorities
loss_for_prior = np.array(elementwise_loss)
new_priorities = loss_for_prior + self.prior_eps
self.memory.update_priorities(indices, new_priorities)

out from the current function and define another new function like:

def update_PER(self,params1,params2,samples1,samples2,indices):
        
        # 1-step Learning loss
        elementwise_loss = self._compute_dqn_loss(params1,params2,samples1, self.gamma)
        
        # N-step Learning loss
        if self.use_n_step:
            gamma = self.gamma ** self.n_step
            elementwise_loss_n_loss = self._compute_dqn_loss(params1, params2, samples2, gamma)
            elementwise_loss += elementwise_loss_n_loss
        
        # PER: update priorities
        loss_for_prior = np.array(elementwise_loss)
        new_priorities = loss_for_prior + self.prior_eps
        self.memory.update_priorities(indices, new_priorities)

And this function is not decorated by the jit and called outside all decorated functions. But then it raises another issue which is I cannot jit or accerelate this function by any JAX decorator, am I right? ><

[jakevdp]

Numpy array conversion aside, it looks like what you're trying to do is fundamentally incompatible with JAX transformations. Your function is impure – each time it is called, it has a side-effect of updating an attribute of self. JAX requires pure functions in order for transforms to operate correctly.

Perhaps you can re-formulate your code so that a copy of the updated object is returned, rather than attempting to mutate the input object within the call?