ConcreteArrays and what to do about them

[shashank2000]

Hey everyone - quick question about JAX/Flax. I’m trying to use the value_and_grad function like so:

def loss_fn(batch):
  return self.child_system.apply(batch, method=self.child_system.training_step)
grad_fn = value_and_grad(loss_fn, has_aux=True)
(loss, acc), grads = grad_fn(batch)        

I looked through the docs, and it seems like not specifying argnums leads to differentiating with respect to the first param by default - which I think shouldn’t be a problem in this use case. In any case, what is happening right now is
batch is a DeviceArray with float32s. The exception raised when I run the code above is AttributeError: 'ConcreteArray' object has no attribute 'items' , and it seems like the ConcreteArray is a result of some internal optimization? I’m not entirely sure why this isn’t working, would greatly appreciate any pointers/things that smell fishy.

EDIT: Turns out I wasn't passing a dict into the first param of apply. This particular error resolved :)

[jakevdp]

Glad you figured it out! In case anyone comes across this question when trying to figure out what JAX ConcreteArrays are, they're basically a type of abstract array representation that JAX uses as a stand-in for an array object during tracing. See Static vs. Traced Operations for a high-level description of how JAX tracing works.