Abstract tracer value encountered where concrete value is expected

[snat1505027]

import dataclasses
import jax
import jax.numpy as jnp
import numpy as np
import flax.linen as nn
from jax import lax
from typing import Any, Dict, Union, Optional
import math
from jax import jit
from copy import deepcopy
import clu.parameter_overview
from dataclasses import dataclass
from functools import partial
import jax.experimental.host_callback as hcb


class VisionTransformer(nn.Module):
    patch_size: int = 16
    hidden_size: int = 768
    size_per_head: int = 64
    dtype: Any = jnp.float32
    num_layers: int = 12
    pooling_ratio: int = 2
    output_grid_h: int = 12
    output_grid_w: int = 20
    do_rotary: bool = True

    @nn.compact
    def __call__(self, x, mask = None):
        """
        :param x: Patched images of size [B, H*W, P * P * 3]
        :return: A pooled representation of size [B, hidden_size]
                 and a seq representation of size [B, HW // (pooling_ratio ** 2), hidden_size]
        """
        *batch_dims, hw, pp3 = x.shape # (7, 540, 768)
        assert hw == (self.output_grid_h * self.output_grid_w) # (18*30)
        assert pp3 == (self.patch_size ** 2) * 3

        x = nn.Dense(features=self.hidden_size, dtype=self.dtype, kernel_init=kernel_init, name='embedding')(x) # (7, 540, 768)

        if self.do_rotary:
            coords = get_rotary_coordinates_2d(self.output_grid_h, self.output_grid_w, dtype=self.dtype) # (540, 2)
        else:
            coords = None
        t_out = TransformerEncoder(hidden_size=self.hidden_size, dtype=self.dtype,
                                   add_cls_token=True, num_layers=self.num_layers, name='transformer',
                                   size_per_head=self.size_per_head)(x, rotary_coords=coords) # 'cls': (7, 768), 'seq': (7, 540, 768) 

        # Attention pool
        assert self.output_grid_h % self.pooling_ratio == 0
        assert self.output_grid_w % self.pooling_ratio == 0
        h2 = self.output_grid_h // self.pooling_ratio # 9
        w2 = self.output_grid_w // self.pooling_ratio # 15
        b2 = int(np.prod(list(batch_dims) + [h2])) # 63

        seq = jnp.reshape(t_out['seq'], [b2, self.pooling_ratio, w2, self.pooling_ratio, self.hidden_size]) # (63,2,15,2,768)
        seq = seq.swapaxes(-4, -3) # (63,15,2,2,768)
        seq = seq.reshape([b2 * w2, self.pooling_ratio ** 2, self.hidden_size]) # (945,4,768)
        target_seq = deepcopy(seq) # (945,4,768)
        target_q = target_seq.mean(-2, keepdims=True) # (945,4,768)
        
        if mask is not None:
            mseq = jnp.reshape(mask, [b2, self.pooling_ratio, w2, self.pooling_ratio, self.hidden_size]) # (63,2,15,2,768)
            mseq = mseq.swapaxes(-4, -3) # (63,15,2,2,768)
            mseq = mseq.reshape([b2 * w2, self.pooling_ratio ** 2, self.hidden_size]) # (945,4,768)
            midx = jnp.unique(jnp.stack([jnp.where(mseq)[:,0], jnp.where(mseq)[:,1]], -1), axis=0)
            

I am trying to get the embedding of image x which will be masked using the tensor mask. So the indexes in which mask contains nonzero entries, the values of those indices in the x will be replaced by zero. But when I call the function I get the following error at this line midx = jnp.unique(jnp.stack([jnp.where(mseq)[:,0], jnp.where(mseq)[:,1]], -1), axis=0):

jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/1)>
The size argument of jnp.nonzero must be statically specified to use jnp.nonzero within JAX transformations.
While tracing the function _microbatch at  /retrain/finetune/tvqa/../../finetune/optimization.py:213 for scan, this concrete value was not available in Python because it depends on the value of the argument 'microbatch'.

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError