def is_module_tree(in_tree: Any) -> bool:   # reject trivial pytrees, {}, [], (), etc.   if not tree_util.tree_leaves(in_tree):     return False-  reduce_fn = lambda prev, cur: prev and isinstance(cur, Module)-  return jax.tree_util.tree_reduce(reduce_fn, in_tree, True)+  reduce_fn = lambda prev, cur: prev or isinstance(cur, Module)

 def is_module_tree(in_tree: Any) -> bool:   # reject trivial pytrees, {}, [], (), etc.

 def f(scope):      init(f)(random.PRNGKey(0)) -  +  def test_inconsistent_param_shapes(self):+    def f(scope):+      scope.param('test', nn.initializers.ones, (4,))+    +    msg = 'Inconisist shapes between value and initializer for paramater "test": (2,), (4,)'

 def __init_subclass__(cls):     cls.scope = None    @classmethod-  def _add_parent_and_name_attrs(cls):-    """Add final optional dataclass attributes: `parent` and `name`."""-    annotations = cls.__dict__.get('__annotations__', {})+  def _customized_dataclass_transform(cls):+    """Handle final optional dataclass attributes: `parent` and `name`."""+    annotations = dict(cls.__dict__.get('__annotations__', {}))     if 'parent' in annotations or 'name' in annotations:       raise ValueError(           f'properties `parent` and `name` are reserved: {annotations}')     # Add `parent` and `name` default fields at end.-    new_annotations = {}-    new_annotations.update(annotations)-    if 'parent' in getattr(cls, '__dataclass_fields__', {}):+    # We temporarily modify base class __dataclass_fields__ to force desired+    # argument behavior and ordering from dataclass class-transform.+    parent_dataclass_fields = dict(getattr(cls, '__dataclass_fields__', {}))+    if 'parent' in parent_dataclass_fields:       cls.__dataclass_fields__.pop('parent')-    new_annotations['parent'] = Union[Type["Module"], Type["Scope"],-                                      Type["_Sentinel"], None]-    cls.parent = dataclasses.field(repr=False, default=_unspecified_parent)-    if 'name' in getattr(cls, '__dataclass_fields__', {}):+    if 'name' in parent_dataclass_fields:       cls.__dataclass_fields__.pop('name')-    new_annotations['name'] = str-    cls.__annotations__ = new_annotations+    annotations['parent'] = Union[Type["Module"], Type["Scope"],+                                  Type["_Sentinel"], None]+    cls.parent = dataclasses.field(repr=False, default=_unspecified_parent)+    annotations['name'] = str     cls.name = None  # default value of name is None.+    cls.__annotations__ = annotations+    # Now apply dataclass transform (which operates in-place).+    dataclasses.dataclass(cls)+    # Restore original base class __dataclass_fields__.+    if dataclasses.is_dataclass(cls.__bases__[0]):+     cls.__bases__[0].__dataclass_fields__ = parent_dataclass_fields

 def is_module_tree(in_tree: Any) -> bool:   # reject trivial pytrees, {}, [], (), etc.   if not tree_util.tree_leaves(in_tree):     return False-  reduce_fn = lambda prev, cur: prev and isinstance(cur, Module)-  return jax.tree_util.tree_reduce(reduce_fn, in_tree, True)+  reduce_fn = lambda prev, cur: prev or isinstance(cur, Module)

 def variable(self, col: str, name: str, init_fn: Callable[..., T],    def param(self, name: str, init_fn: Callable[..., T], *init_args) -> T:     """Create a paramater."""-    s_init_fn = lambda *args: init_fn(self.make_rng('params'), *init_args)-    v = self.variable('params', name, s_init_fn, *init_args)-    return v.value+    if self.has_variable('params', name):+      abs_rng = jax.ShapeDtypeStruct((2,), jnp.uint32)+      value = self.get_variable('params', name)+      # validate shape of init_fn output is the same as the shape of the existing+      # paramater.+      abs_value = jax.eval_shape(lambda rng: init_fn(rng, *init_args), abs_rng)+      abs_value_flat = jax.tree_leaves(abs_value)+      value_flat = jax.tree_leaves(value)+      for val, abs_val in zip(value_flat, abs_value_flat):+        # NOTE: we could check dtype consistency here as well but it's usefuleness is less obvious.+        # we might intentionally change the dtype for inference to a half float type for example.+        if jnp.shape(val) != jnp.shape(abs_val):+          raise ValueError('Inconisist shapes between value and initializer '+                           f'for paramater "{name}": {jnp.shape(val)}, {jnp.shape(abs_val)}')

 def variable(self, col: str, name: str, init_fn: Callable[..., T],    def param(self, name: str, init_fn: Callable[..., T], *init_args) -> T:     """Create a paramater."""-    s_init_fn = lambda *args: init_fn(self.make_rng('params'), *init_args)-    v = self.variable('params', name, s_init_fn, *init_args)-    return v.value+    if self.has_variable('params', name):+      abs_rng = jax.ShapeDtypeStruct((2,), jnp.uint32)+      value = self.get_variable('params', name)+      # validate shape of init_fn output is the same as the shape of the existing+      # paramater.+      abs_value = jax.eval_shape(lambda rng: init_fn(rng, *init_args), abs_rng)+      abs_value_flat = jax.tree_leaves(abs_value)+      value_flat = jax.tree_leaves(value)+      for val, abs_val in zip(value_flat, abs_value_flat):+        # NOTE: we could check dtype consistency here as well but it's usefuleness is less obvious.+        # we might intentionally change the dtype for inference to a half float type for example.+        if jnp.shape(val) != jnp.shape(abs_val):+          raise ValueError('Inconisist shapes between value and initializer '

 def variable(self, col: str, name: str, init_fn: Callable[..., T],    def param(self, name: str, init_fn: Callable[..., T], *init_args) -> T:     """Create a paramater."""-    s_init_fn = lambda *args: init_fn(self.make_rng('params'), *init_args)-    v = self.variable('params', name, s_init_fn, *init_args)-    return v.value+    if self.has_variable('params', name):+      abs_rng = jax.ShapeDtypeStruct((2,), jnp.uint32)+      value = self.get_variable('params', name)+      # validate shape of init_fn output is the same as the shape of the existing+      # paramater.+      abs_value = jax.eval_shape(lambda rng: init_fn(rng, *init_args), abs_rng)+      abs_value_flat = jax.tree_leaves(abs_value)+      value_flat = jax.tree_leaves(value)+      for val, abs_val in zip(value_flat, abs_value_flat):+        # NOTE: we could check dtype consistency here as well but it's usefuleness is less obvious.+        # we might intentionally change the dtype for inference to a half float type for example.+        if jnp.shape(val) != jnp.shape(abs_val):+          raise ValueError('Inconisist shapes between value and initializer '

 def __call__(self, carry, x):     carry_rng, categorical_rng = jax.random.split(rng, 2)     if not self.teacher_force:       x = last_prediction-    lstm_cell = nn.LSTMCell(name='lstm_cell')-    projection = nn.Dense(features=self.vocab_size, name='projection')-    lstm_state, y = lstm_cell(lstm_state, x)-    logits = projection(y)-    predicted_tokens = jax.random.categorical(categorical_rng, logits)-    prediction = onehot(predicted_tokens, self.vocab_size)+    lstm_state, y = nn.LSTMCell()(lstm_state, x)+    logits = nn.Dense(features=CTABLE.vocab_size)(y)+    predicted_token = jax.random.categorical(categorical_rng, logits)+    prediction = jnp.array(predicted_token == jnp.arange(CTABLE.vocab_size), +                           dtype=jnp.float32)

 def compute_metrics(logits, labels):   return metrics  +IN_SHAPES = [{'query': '(n, _)', 'answer': '(m, _)'}]+OUT_ELEM = f'(m +- 1, {CTABLE.vocab_size})'

 def encode_onehot(batch_inputs, max_len):    def encode_str(s):     tokens = CTABLE.encode(s)-    if len(tokens) > max_len:+    unpadded_len = len(tokens)+    if unpadded_len > max_len:       raise ValueError(f'Sequence too long ({len(tokens)}>{max_len}): \'{s}\'')     tokens = np.pad(tokens, [(0, max_len-len(tokens))], mode='constant')-    return onehot(tokens, CTABLE.vocab_size)+    return onehot(tokens, CTABLE.vocab_size), unpadded_len -  return np.array([encode_str(inp) for inp in batch_inputs])+  return np.array([encode_str(inp) for inp in batch_inputs], dtype=object)

 def __call__(self, carry, x):     carry_rng, categorical_rng = jax.random.split(rng, 2)     if not self.teacher_force:       x = last_prediction-    lstm_cell = nn.LSTMCell(name='lstm_cell')-    projection = nn.Dense(features=self.vocab_size, name='projection')-    lstm_state, y = lstm_cell(lstm_state, x)-    logits = projection(y)-    predicted_tokens = jax.random.categorical(categorical_rng, logits)-    prediction = onehot(predicted_tokens, self.vocab_size)+    lstm_state, y = nn.LSTMCell()(lstm_state, x)+    logits = nn.Dense(features=CTABLE.vocab_size)(y)+    predicted_token = jax.random.categorical(categorical_rng, logits)+    prediction = jnp.array(predicted_token == jnp.arange(CTABLE.vocab_size), +                           dtype=jnp.float32)

 def encode_onehot(batch_inputs, max_len):    def encode_str(s):     tokens = CTABLE.encode(s)-    if len(tokens) > max_len:+    unpadded_len = len(tokens)+    if unpadded_len > max_len:       raise ValueError(f'Sequence too long ({len(tokens)}>{max_len}): \'{s}\'')     tokens = np.pad(tokens, [(0, max_len-len(tokens))], mode='constant')-    return onehot(tokens, CTABLE.vocab_size)+    return onehot(tokens, CTABLE.vocab_size), unpadded_len -  return np.array([encode_str(inp) for inp in batch_inputs])+  return np.array([encode_str(inp) for inp in batch_inputs], dtype=object)

 def compute_metrics(logits, labels):   return metrics  +IN_SHAPES = [{'query': '(n, _)', 'answer': '(m, _)'}]+OUT_ELEM = f'(m +- 1, {CTABLE.vocab_size})'

+from typing import Callable, Dict++import numpy as np++import flax.nn as nn+import jax+import jax.numpy as jnp+from transformers import BertConfig+from transformers.modeling_bert import BERT_PRETRAINED_MODEL_ARCHIVE_MAP+from transformers.modeling_flax_utils import FlaxPreTrainedModel++[email protected]+def gelu(x):+    r"""Gaussian error linear unit activation function.++    Computes the element-wise function:++    .. math::+      \mathrm{gelu}(x) = \frac{x}{2} \left(1 + \mathrm{tanh} \left(+        \sqrt{\frac{2}{\pi}} \left(x + 0.044715 x^3 \right) \right) \right)++    We explicitly use the approximation rather than the exact formulation for+    speed. For more information, see `Gaussian Error Linear Units (GELUs)+    <https://arxiv.org/abs/1606.08415>`_, section 2.+    """+    return x * 0.5 * (1.0 + jax.lax.erf(x / jnp.sqrt(2.0)))+++ACT2FN = {+    "gelu": gelu,+    "relu": nn.relu,+    "swish": nn.swish,+    "gelu_new": gelu,+}+++class BertLayerNorm(nn.Module):+    """Layer normalization (https://arxiv.org/abs/1607.06450).+    Operates on the last axis of the input data.+    """++    def apply(+        self,+        x,+        epsilon=1e-6,+        dtype=jnp.float32,+        bias=True,+        scale=True,+        bias_init=nn.initializers.zeros,+        scale_init=nn.initializers.ones,+    ):+        """Applies layer normalization on the input.+        It normalizes the activations of the layer for each given example in a+        batch independently, rather than across a batch like Batch Normalization.+        i.e. applies a transformation that maintains the mean activation within+        each example close to 0 and the activation standard deviation close to 1.+        Args:+          x: the inputs+          epsilon: A small float added to variance to avoid dividing by zero.+          dtype: the dtype of the computation (default: float32).+          bias:  If True, bias (beta) is added.+          scale: If True, multiply by scale (gamma). When the next layer is linear+            (also e.g. nn.relu), this can be disabled since the scaling will be done+            by the next layer.+          bias_init: Initializer for bias, by default, zero.+          scale_init: Initializer for scale, by default, one.+        Returns:+          Normalized inputs (the same shape as inputs).+        """+        features = x.shape[-1]+        mean = jnp.mean(x, axis=-1, keepdims=True)+        mean2 = jnp.mean(jnp.lax.square(x), axis=-1, keepdims=True)+        var = mean2 - jnp.lax.square(mean)+        mul = jnp.lax.rsqrt(var + epsilon)+        if scale:+            mul = mul * jnp.asarray(self.param("gamma", (features,), scale_init), dtype)+        y = (x - mean) * mul+        if bias:+            y = y + jnp.asarray(self.param("beta", (features,), bias_init), dtype)+        return y+++class BertEmbedding(nn.Module):+    """+    Specify a new class for doing the embedding stuff+    as Flax's one use 'embedding' for the parameter name+    and PyTorch use 'weight'+    """++    def apply(+        self,+        input,+        vocab_size: int,+        hidden_size: int,+        emb_init: Callable[..., np.ndarray] = nn.initializers.normal(stddev=0.1),+    ):++        embedding = self.param("weight", (vocab_size, hidden_size), emb_init)+        return jnp.take(embedding, input, axis=0)+++class BertEmbeddings(nn.Module):+    def apply(+        self,+        input_ids,+        token_type_ids,+        position_ids,+        attention_mask,+        vocab_size: int,+        hidden_size: int,+        type_vocab_size: int,+        max_length: int,+    ):++        # Embed+        w_emb = BertEmbedding(jnp.atleast_2d(input_ids.astype("i4")), vocab_size, hidden_size, name="word_embeddings")+        p_emb = BertEmbedding(+            jnp.atleast_2d(position_ids.astype("i4")), max_length, hidden_size, name="position_embeddings"+        )+        t_emb = BertEmbedding(+            jnp.atleast_2d(token_type_ids.astype("i4")), type_vocab_size, hidden_size, name="token_type_embeddings"+        )++        # Sum all embeddings+        summed_emb = w_emb + jnp.broadcast_to(p_emb, w_emb.shape) + t_emb++        # Layer Norm+        norm = BertLayerNorm(summed_emb, name="layer_norm")++        return norm+++class BertAttention(nn.Module):+    def apply(self, hidden_state, attention_mask, num_heads: int, head_size: int):+        self_att = nn.attention.SelfAttention(+            hidden_state, num_heads=num_heads, qkv_features=head_size, padding_mask=attention_mask, name="self"+        )++        return BertLayerNorm(self_att + hidden_state, name="layer_norm")+++class BertIntermediate(nn.Module):+    def apply(self, hidden_state, output_size: int):+        # TODO: Add ACT2FN reference to change activation function+        h = nn.Dense(hidden_state, features=output_size, name="dense")+        return gelu(h)+++class BertOutput(nn.Module):+    def apply(self, intermediate_output, attention_output):+        h = nn.Dense(intermediate_output, attention_output.shape[-1], name="dense")+        h = BertLayerNorm(h + attention_output, name="layer_norm")++        return h+++class BertLayer(nn.Module):+    def apply(self, hidden_state, attention_mask, num_heads: int, head_size: int, intermediate_size: int):+        attention = BertAttention(hidden_state, attention_mask, num_heads, head_size, name="attention")+        intermediate = BertIntermediate(attention, intermediate_size, name="intermediate")+        output = BertOutput(intermediate, attention, name="output")++        return output+++class BertLayerCollection(nn.Module):+    """+    Stores N BertLayer(s)+    """++    def apply(self, input, attention_mask, num_layers: int, num_heads: int, head_size: int, intermediate_size: int):+        assert num_layers > 0, "num_layers should be >= 1, got ({})".format(num_layers)++        # Initialize input / output+        input_i = output_i = input++        # Forward over all encoders+        for i in range(num_layers):+            output_i = BertLayer(input_i, attention_mask, num_heads, head_size, intermediate_size, name="{}".format(i))+            input_i = output_i+        return output_i+++class BertEncoder(nn.Module):+    def apply(+        self, hidden_state, attention_mask, num_layers: int, num_heads: int, head_size: int, intermediate_size: int+    ):+        encodings = BertLayerCollection(+            hidden_state, attention_mask, num_layers, num_heads, head_size, intermediate_size, name="layer"+        )+        return encodings+++class BertPooler(nn.Module):+    def apply(self, hidden_state):+        first_token = hidden_state[:, 0]+        out = nn.Dense(first_token, hidden_state.shape[-1], name="dense")+        return jax.lax.tanh(out)+++class BertModel(nn.Module):+    def apply(+        self,+        input_ids,+        token_type_ids,+        position_ids,+        attention_mask,+        vocab_size: int,+        hidden_size: int,+        type_vocab_size: int,+        max_length: int,+        num_encoder_layers: int,+        num_heads: int,+        head_size: int,+        intermediate_size: int,+        padding_idx: int,+        emb_init: Callable[..., np.ndarray] = nn.initializers.normal(stddev=0.1),+    ):++        # Embedding+        embeddings = BertEmbeddings(+            input_ids,+            token_type_ids,+            position_ids,+            attention_mask,+            vocab_size,+            hidden_size,+            type_vocab_size,+            max_length,+            name="embeddings",+        )++        # N stacked encoding layers+        encoder = BertEncoder(+            embeddings, attention_mask, num_encoder_layers, num_heads, head_size, intermediate_size, name="encoder"+        )++        pooled = BertPooler(encoder, name="pooler")+        return encoder, pooled+++class FlaxBertModel(FlaxPreTrainedModel):+    """+    BERT implementation using JAX/Flax as backend+    """++    model_class = BertModel+    config_class = BertConfig+    pretrained_model_archive_map = BERT_PRETRAINED_MODEL_ARCHIVE_MAP+    base_model_prefix = "bert"++    @staticmethod+    def convert_from_pytorch(pt_state: Dict, config: BertConfig) -> Dict:+        jax_state = dict(pt_state)++        # Need to change some parameters name to match Flax names so that we don't have to fork any layer+        for key, tensor in pt_state.items():+            # Key parts+            key_parts = set(key.split("."))++            # Every dense layer have a "kernel" parameters instead of "weight"+            if "dense.weight" in key:+                del jax_state[key]+                key = key.replace("weight", "kernel")+                jax_state[key] = tensor++            # SelfAttention needs also to replace "weight" by "kernel"+            if {"query", "key", "value"} & key_parts:++                # Flax SelfAttention decomposes the heads (num_head, size // num_heads)+                if "bias" in key:+                    jax_state[key] = tensor.reshape((config.num_attention_heads, -1))+                elif "weight":+                    del jax_state[key]+                    key = key.replace("weight", "kernel")+                    tensor = tensor.reshape((config.num_attention_heads, -1, config.hidden_size)).transpose((2, 0, 1))+                    jax_state[key] = tensor++            # SelfAttention output is not a separate layer, remove one nesting+            if "attention.output.dense" in key:+                del jax_state[key]+                key = key.replace("attention.output.dense", "attention.self.out")+                jax_state[key] = tensor++            # SelfAttention output is not a separate layer, remove nesting on layer norm+            if "attention.output.LayerNorm" in key:+                del jax_state[key]+                key = key.replace("attention.output.LayerNorm", "attention.LayerNorm")+                jax_state[key] = tensor++            # There are some transposed parameters w.r.t their PyTorch counterpart+            if "intermediate.dense.kernel" in key or "output.dense.kernel" in key:+                jax_state[key] = tensor.T++            # Self Attention output projection needs to be transposed+            if "out.kernel" in key:+                jax_state[key] = tensor.reshape((config.hidden_size, config.num_attention_heads, -1)).transpose(+                    1, 2, 0+                )++            # Pooler needs to transpose its kernel+            if "pooler.dense.kernel" in key:+                jax_state[key] = tensor.T++            # Handle LayerNorm conversion+            if "LayerNorm" in key:+                del jax_state[key]++                # Replace LayerNorm by layer_norm+                new_key = key.replace("LayerNorm", "layer_norm")++                if "weight" in key:+                    new_key = new_key.replace("weight", "gamma")+                elif "bias" in key:+                    new_key = new_key.replace("bias", "beta")++                jax_state[new_key] = tensor++        return jax_state++    def __init__(self, config: BertConfig, state: dict, seed: int = 0, **kwargs):+        model_def = BertModel.partial(+            vocab_size=config.vocab_size,+            hidden_size=config.hidden_size,+            type_vocab_size=config.type_vocab_size,+            max_length=config.max_position_embeddings,+            num_encoder_layers=config.num_hidden_layers,+            num_heads=config.num_attention_heads,+            head_size=config.hidden_size,+            intermediate_size=config.intermediate_size,+            padding_idx=config.pad_token_id,+        )++        super().__init__(config, model_def, state, seed)++    @property+    def module(self) -> BertModel:+        return self._module++    @property+    def config(self) -> BertConfig:+        return self._config++    def __call__(self, input_ids, token_type_ids=None, position_ids=None, attention_mask=None):+        @jax.jit

     _tf_available = False  # pylint: disable=invalid-name  +try:+    USE_JAX = os.environ.get("USE_FLAX", "AUTO").upper()++    if USE_JAX in ENV_VARS_TRUE_VALUES:+        import flax+        import jax+        from jax.config import config+        # TODO(marcvanzee): Flax Linen requires JAX omnistaging. Remove this +        # once JAX enables it by default.+        config.enable_omnistaging()

 def find_length(axis, x):       return ()     # split rngs     lengths = jax.tree_multimap(find_length, in_axes, args)-    if length is None:-      d_length, = set(jax.tree_leaves(lengths))+    lengths = set(jax.tree_leaves(lengths))+    if length is None and len(lengths) == 1:+      d_length, = lengths+    elif len(lengths) > 1:+      raise ValueError(f'Incosistent scan lengths: {lengths}')

  import numpy as onp +import tensorflow as tf+tf.config.experimental.set_visible_devices([], "GPU")

  import numpy as onp +import tensorflow as tf+tf.config.experimental.set_visible_devices([], "GPU")

 import jax.numpy as jnp  import tensorflow as tf+tf.config.experimental.set_visible_devices([], "GPU")

  PRNGKey = Any Shape = Tuple[int]-Dtype = Any  # this could be a real type?+Dtype = Any Array = Any  -def dot_product_attention(query,-                          key,-                          value,-                          dtype=jnp.float32,-                          bias=None,-                          axis=None,-                          broadcast_dropout=True,-                          dropout_rng=None,-                          dropout_rate=0.,-                          deterministic=False,-                          precision=None):+def dot_product_attention(query: Array,+                          key: Array,+                          value: Array,+                          bias: Optional[Array] = None,+                          broadcast_dropout: bool = True,

 def __call__(self,         query,         key,         value,-        dtype=self.dtype,-        axis=attention_axis,         bias=attention_bias,-        precision=self.precision,         dropout_rng=dropout_rng,         dropout_rate=self.dropout_rate,         broadcast_dropout=self.broadcast_dropout,-        deterministic=self.deterministic)+        deterministic=self.deterministic,+        dtype=self.dtype,+        precision=self.precision)      # back to the original inputs dimensions     out = DenseGeneral(features=features,

  ScanVariableMode = Union[str, Tuple[str, str]] ++def _dedup_scopes(scopes):+  paths = []+  minimal_set = set(scopes)+  for leave in scopes: