Add files using upload-large-folder tool

Tipsomaly/model/big_vision/configs/proj/paligemma/transfers/tallyqa.py

@@ -0,0 +1,191 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to TallyQA.
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+
+def training_data(res, text_len=32):
+  """Creates training data config.
+
+  See (internal link)
+  You can add more arguments beside `res`, but give them good defaults.
+
+  Args:
+    res: The requested image resolution (eg 224)
+    text_len: sequence length
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  c = bvcc.parse_arg('')  # Just make a configdict without extra import.
+  c.data = dict(
+      name='tallyqa',
+      split='train',
+  )
+
+  c.pp = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      'strfmt("answer en {question}", outkey="prefix")',
+      'strfmt("{answer}", outkey="suffix")',
+      combine_and_keep_train(text_len),
+  ])
+  return c
+
+
+def countbenchqa_eval_data(res, text_len=32):
+  """Creates eval data config for CountBenchQA."""
+  c = bvcc.parse_arg('')  # Just make a configdict without extra import.
+  c.data = dict(
+      name='countbenchqa',
+      split='huggingface',
+  )
+  c.pp = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      'strfmt("answer en {question}", outkey="prefix")',
+      'strfmt("{number}", outkey="answer")',
+      combine_and_keep_eval(text_len, keep=('answer',)),
+  ])
+  return c
+
+
+def add_eval(c, res, text_len=32, **kw):
+  """Add eval configs."""
+  tallyqa_pp_eval = '|'.join([
+      f'decode|resize({res}, antialias=True)|value_range(-1, 1)',
+      'strfmt("answer en {question}", outkey="prefix")',
+      'strfmt("{answer}", outkey="answer")',
+      combine_and_keep_eval(text_len, keep=('answer', 'issimple')),
+  ])
+
+  for freq, name, split in [
+      (0.1, 'minitrain', 'train[:5%]'),
+      # (0.1, 'minival', 'train[-5%:]'),
+      (1/4, 'eval', 'test'),
+  ]:
+    c.evals[f'tallyqa/{name}'] = dict(
+        type='proj.paligemma.transfers.tallyqa',
+        pred='decode', pred_kw={'max_decode_len': text_len},
+        data={**training_data(res, text_len).data, 'split': split},
+        log_percent=freq, tokenizer=TOKENIZER, pp_fn=tallyqa_pp_eval)
+    c.evals[f'tallyqa/{name}'].update(kw)
+
+  # CountBenchQA eval. We use the TallyQA eval for this but just pass in
+  # different data.
+  c.evals['countbenchqa/eval'] = dict(
+      type='proj.paligemma.transfers.tallyqa',
+      pred='decode', pred_kw={'max_decode_len': text_len},
+      data=countbenchqa_eval_data(res, text_len).data,
+      log_percent=0.1,  # This is a very small and cheap eval set.
+      tokenizer=TOKENIZER,
+      pp_fn=countbenchqa_eval_data(res, text_len).pp)
+  c.evals['countbenchqa/eval'].update(kw)
+
+
+def add_eval_pplx(c, res, text_len=32):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, text_len)  # Use mostly same settings as training.
+  for name, split in [
+      ('minitrain', 'train[:5%]'),  # To gauge memorization.
+      # ('minival', 'train[-5%:]'),  # To tune hparams.
+      ('eval', 'test'),  # To compute final publishable scores.
+  ]:
+    c.evals[f'tallyqa/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=0.1,  # Eval ~10x per run; it's cheap.
+        data={**c_train.data, 'split': split},
+        pp_fn=c_train.pp,
+    )
+
+
+def sweep_best(add, arg=None):  # pylint: disable=unused-argument
+  """Train with best hyper-params."""
+  add(total_epochs=2, lr=1e-5, wd=0.00, **bvcc.arg(res=224))
+  add(total_epochs=2, lr=1e-5, wd=1e-6, **bvcc.arg(res=448))
+  add(total_epochs=2, lr=7e-6, wd=7e-7, **bvcc.arg(res=896))
+
+
+sweep = sweep_best  # Choose which sweep to run.
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', res=224)
+
+  c.input = training_data(c.res)
+
+  # Instead of epochs, you can also use `total_examples` or `total_steps`.
+  c.total_epochs = 2
+  c.input.batch_size = 256
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 1e-5
+  c.wd = 0.0
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.0
+  c.schedule = dict(decay_type='cosine', warmup_percent=0.05)
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval(c, c.res, batch_size=256)
+  add_eval_pplx(c, c.res)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.0)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  # These probably do not need any change/tuning
+  c.input.shuffle_buffer_size = 50_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'proj.paligemma.ops']
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    c.input.shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.split = ev.data.split.split('[')[0] + '[:16]'
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  """Returns a list of metric names."""
+  return [
+      'training_loss',
+      'countbenchqa/eval/acc',
+      'tallyqa/minitrain/pplx/avg',
+      'tallyqa/eval/pplx/avg',
+      'tallyqa/eval/acc',
+      'tallyqa/eval/acc/complex',
+      'tallyqa/eval/acc/simple',
+  ]

Tipsomaly/model/big_vision/configs/proj/paligemma/transfers/textcaps.py

@@ -0,0 +1,181 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=line-too-long
+r"""PaliGemma transfer to TextCaps captioning task.
+"""
+
+import big_vision.configs.common as bvcc
+from big_vision.configs.proj.paligemma.transfers.common import combine_and_keep_train, combine_and_keep_eval, TOKENIZER
+
+
+def training_data(res, *, final_split, text_len=32, crop='rs'):
+  """Creates training data config.
+
+  See (internal link)
+  You can add more arguments beside `res`, but give them good defaults.
+
+  Args:
+    res: The requested image resolution (eg 224).
+    final_split: Train on all train data or train[:98%].
+    text_len: sequence length.
+    crop: one of {'ic', 'rc', 'rs'}.
+
+  Returns:
+    The ConfigDict for the input section.
+  """
+  c = bvcc.parse_arg('')  # Just make a configdict without extra import.
+  c.data = dict(
+      name='text_caps',
+      split='train' if final_split else 'train[:98%]',
+  )
+
+  if crop == 'ic':
+    crop = f'inception_crop({res}, area_min=50)'
+  elif crop == 'rc':
+    crop = f'resize_small({res*8//7})|random_crop({res})'
+  elif crop == 'rs':
+    crop = f'resize({res})'
+  else:
+    raise ValueError(f'Unknown crop: {crop}')
+
+  c.pp = '|'.join([
+      'flatten',
+      'decode', crop, 'value_range(-1, 1)',
+      'choice_no_replacement(inkey="texts", outkey="suffix")',
+      'strfmt("caption en", outkey="prefix")',
+      'lower(key="suffix")',
+      combine_and_keep_train(text_len),
+  ])
+  return c
+
+
+def add_eval(c, res, text_len=32, **kw):
+  """Captioning evaluator with cider/bleu-4/meteor/rouge/spice metrics."""
+  # Input eval pp without ground truth text and random crop.
+  pp_eval = '|'.join([
+      'decode', f'resize({res})', 'value_range(-1, 1)',
+      'flatten', 'copy("texts", "captions")',  # GT for evaluator.
+      'strfmt("caption en", outkey="prefix")',
+      combine_and_keep_eval(text_len, keep=('image/id', 'captions')),
+  ])
+
+  for name, split in [
+      ('minitrain', 'train[:2%]'),
+      ('minival', 'train[-2%:]'),
+      ('eval', 'val'),
+  ]:
+    c.evals[f'textcaps/{name}'] = dict(
+        type='proj.paligemma.transfers.coco_caption',
+        pred='decode', pred_kw={'max_decode_len': text_len},
+        log_percent=0.1, tokenizer=TOKENIZER,
+        data={'name': 'text_caps', 'split': split},
+        pp_fn=pp_eval,
+    )
+    c.evals[f'textcaps/{name}'].update(kw)
+
+
+def add_eval_pplx(c, res, text_len=32):
+  """Perplexity evaluator to test runs before implementing the real deal."""
+  c_train = training_data(res, final_split=True, text_len=text_len, crop='rs')  # Use mostly same settings as training.
+  for name, split in [
+      ('minitrain', 'train[:2%]'),
+      ('minival', 'train[-2%:]'),
+      ('eval', 'val'),
+  ]:
+    c.evals[f'textcaps/{name}/pplx'] = dict(
+        type='proj.paligemma.perplexity', pred='logits',
+        key='text', shift_labels=True,
+        log_percent=0.05,  # Eval ~20x per run; it's cheap.
+        data={**c_train.data, 'split': split},
+        pp_fn=c_train.pp,
+    )
+
+
+def get_config(arg=None):
+  """Config for training."""
+  c = bvcc.parse_arg(arg, mode='xm', crop='rs', res=224, beam_size=3, final_split=False)
+
+  c.input = training_data(c.res, final_split=c.final_split, crop=c.crop)
+
+  c.total_epochs = 5  # Note each example has 5 captions.
+  c.input.batch_size = 256
+  c.optax_name = 'scale_by_adam'
+  c.optax = dict(b2=0.999)
+  c.lr = 1e-5
+  c.wd = 0.0
+  c.grad_clip_norm = 1.0
+  c.label_smoothing = 0.0
+  c.schedule = dict(decay_type='cosine', warmup_percent=0.05)
+
+  # Add evaluators.
+  c.evals = {}
+  add_eval_pplx(c, c.res)
+
+  if c.beam_size:
+    decode_kw = {'pred': 'beam_decode', 'pred_kw': {'beam_size': c.beam_size}}
+  else:
+    decode_kw = {}
+
+  add_eval(c, c.res, batch_size=1024, **decode_kw)
+
+  # Model section.
+  c.model_name = 'proj.paligemma.paligemma'
+  c.model = {}
+  c.model.img = dict(variant='So400m/14', pool_type='none', scan=True)
+  c.model.llm = dict(vocab_size=256_000 + 1024 + 128, dropout=0.0)
+  c.model_init = f'pt_{c.res}'
+
+  # FSDP strategy.
+  c.mesh = [('data', -1)]
+  c.sharding_strategy = [('.*', 'fsdp(axis="data")')]
+  c.sharding_rules = [('act_batch', ('data',))]
+
+  # These probably do not need any change/tuning
+  c.input.shuffle_buffer_size = 50_000
+  c.log_training_steps = 50
+  c.ckpt_steps = 1_000
+  c.pp_modules = ['ops_general', 'ops_image', 'ops_text', 'proj.paligemma.ops']
+
+  # Update configs for quicker local runs and avoid swapping.
+  if c.mode in ('runlocal', 'mock'):
+    c.input.shuffle_buffer_size = None
+    for ev in c.evals.values():
+      ev.data.split = ev.data.split.split('[')[0] + '[:16]'
+
+  if c.mode == 'runlocal':
+    c.log_training_steps = 1
+    c.input.batch_size = 2
+
+  c.seed = 0
+  return c
+
+
+def sweep_best(add, arg=None):
+  """Train with best hyper-params."""
+  c = bvcc.parse_arg(arg, final_split=False)
+  # Note: wd=0.0 probably works as good.
+  add(lr=1e-5, wd=1e-6, total_epochs=5, **bvcc.arg(res=224, **c))
+  add(lr=1e-5, wd=1e-6, total_epochs=5, **bvcc.arg(res=448, **c))
+
+
+sweep = sweep_best  # Choose which sweep to run.
+
+
+def metrics(arg=None):  # pylint: disable=unused-argument
+  m = ['training_loss']
+  for split in ('eval', 'minival', 'minitrain'):
+    m.append(('epoch', f'textcaps/{split}/cider'))
+    m.append(('epoch', f'textcaps/{split}/pplx/avg'))
+  return m

Tipsomaly/model/big_vision/models/ppp/gemma.py

@@ -0,0 +1,651 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""gemma reimplementation for big_vision.
+
+We follow this einsum axis naming convention:
+  B: batch
+  T: query length
+  S: k/v length
+  N: num query heads
+  K: num k/v heads
+  G: num query heads per k/v head
+  H: head dim
+  D: d_model ("features")
+
+Example Colab using the models via the PaliGemma decoding logic:
+(internal link)
+
+Doc locating the variable initializers in the original code and validating them:
+(internal link)
+
+This implementation does *not* currently support the local sliding attention
+pattern used in the v2 models. But since we mostly use sequences <4096 tokens,
+this shouldn't make any difference. Since RoPE embedding is used throughout,
+it's unclear if there is any practical difference (other than wasting some
+memory).
+"""
+
+
+from big_vision.models import common
+import big_vision.utils as u
+import einops
+import flax
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import ml_collections
+import numpy as np
+import orbax.checkpoint
+
+
+def get_config(variant):
+  """Returns config for specified gemma variant."""
+  if variant == "gemma_2b":
+    return ml_collections.ConfigDict(
+        dict(
+            variant=variant,
+            width=2048,
+            depth=18,
+            mlp_dim=16_384,
+            num_heads=8,
+            num_kv_heads=1,
+            head_dim=256,
+            norm_eps=1e-6,
+            vocab_size=256_000,
+            scan=True,
+            remat_policy="nothing_saveable",
+        )
+    )
+  if variant == "gemma_7b":
+    return ml_collections.ConfigDict(
+        dict(
+            variant=variant,
+            width=3072,
+            depth=28,
+            mlp_dim=24_576,
+            num_heads=16,
+            num_kv_heads=16,
+            head_dim=256,
+            norm_eps=1e-6,
+            vocab_size=256_000,
+            scan=True,
+            remat_policy="nothing_saveable",
+        )
+    )
+  if variant == "gemma2_2b":
+    return ml_collections.ConfigDict(
+        dict(
+            variant=variant,
+            width=2304,
+            depth=26,
+            mlp_dim=9_216,
+            num_heads=8,
+            num_kv_heads=4,
+            head_dim=256,
+            norm_eps=1e-6,
+            vocab_size=256_000,
+            final_logits_softcap=30.0,
+            attn_logits_softcap=50.0,
+            post_norms=True,
+            scan=True,
+            remat_policy="nothing_saveable",
+        )
+    )
+  if variant == "gemma2_9b":
+    return ml_collections.ConfigDict(
+        dict(
+            variant=variant,
+            width=3584,
+            depth=42,
+            mlp_dim=14_336,
+            num_heads=16,
+            num_kv_heads=8,
+            head_dim=256,
+            norm_eps=1e-6,
+            vocab_size=256_000,
+            final_logits_softcap=30.0,
+            attn_logits_softcap=50.0,
+            post_norms=True,
+            scan=True,
+            remat_policy="nothing_saveable",
+        )
+    )
+  if variant == "gemma2_27b":
+    return ml_collections.ConfigDict(
+        dict(
+            variant=variant,
+            width=4608,
+            depth=46,
+            mlp_dim=36_864,
+            num_heads=32,
+            num_kv_heads=16,
+            head_dim=128,
+            norm_eps=1e-6,
+            vocab_size=256_000,
+            query_pre_attn_norm="rsqrt_emb_per_head",
+            final_logits_softcap=30.0,
+            attn_logits_softcap=50.0,
+            post_norms=True,
+            scan=True,
+            remat_policy="nothing_saveable",
+        )
+    )
+  raise ValueError(f"Unknown variant: {variant}")
+
+
+def _apply_rope(x, *, positions, max_wavelength=10_000):
+  """Applies RoPE positions [B, L] to x [B, L, H, D]."""
+  freq_exponents = (2. / x.shape[-1]) * jnp.arange(x.shape[-1] // 2)
+  timescale = (max_wavelength ** freq_exponents)
+  radians = positions[..., None] / timescale[None, None, :]
+  radians = radians[..., None, :]
+  # radians.shape = [...,L,1,d=D/2]
+  sin, cos = jnp.sin(radians), jnp.cos(radians)
+  x1, x2 = jnp.split(x, 2, axis=-1)
+  res = jnp.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1)
+  return res
+
+
+def _update_kv_cache(module, k, v, cache_size, cache_dtype):
+  """Updates KV cache and returns its current contents."""
+  initialized = module.has_variable("cache", "idx")
+  batch_size, update_len, num_heads, head_dim = k.shape
+  cache_dtype = cache_dtype or k.dtype
+
+  # Idx of which cache row to update next is the same for all examples, so that
+  # it allows to update with dynamic_update_slice. But in order to keep things
+  # nicely partitioned we store it with leading batch dimension and use only
+  # the first entry.
+  idx = module.variable("cache", "idx", jnp.zeros, (batch_size,), jnp.int32)
+
+  kv_shape = (batch_size, cache_size, num_heads, head_dim)
+  k_cache = module.variable(
+      "cache", "k_cache", jnp.zeros, kv_shape, cache_dtype)
+  v_cache = module.variable(
+      "cache", "v_cache", jnp.zeros, kv_shape, cache_dtype)
+
+  if initialized:  # write k, v in the next cache position.
+    assert update_len == 1, update_len
+    # Note: idx is the same for all examples. Use value from example 0.
+    indices = (0, idx.value[0], 0, 0)
+    k_cache.value = jax.lax.dynamic_update_slice(
+        k_cache.value, k.astype(cache_dtype), indices)
+    v_cache.value = jax.lax.dynamic_update_slice(
+        v_cache.value, v.astype(cache_dtype), indices)
+    idx.value = idx.value + 1
+  else:  # init cache with k, v after padding to cache_size.
+    prefill_len = k.shape[1]
+    pad_width = ((0, 0), (0, cache_size - prefill_len), (0, 0), (0, 0))
+    k_cache.value = jnp.pad(k.astype(cache_dtype), pad_width)
+    v_cache.value = jnp.pad(v.astype(cache_dtype), pad_width)
+    idx.value = idx.value + prefill_len
+
+  return k_cache.value.astype(k.dtype), v_cache.value.astype(v.dtype)
+
+
+def trunc_norm_init(in_axis, out_axis, batch_axis):
+  return nn.initializers.variance_scaling(
+      1.0, "fan_in", "truncated_normal",
+      in_axis=in_axis, out_axis=out_axis, batch_axis=batch_axis)
+
+
+class Einsum(nn.Module):
+  shape: tuple[int, ...]
+  w_init: nn.initializers.Initializer = nn.initializers.zeros_init()
+
+  @nn.compact
+  def __call__(self, eqn, x):
+    w = self.param("w", self.w_init, self.shape)
+    return jnp.einsum(eqn, x, w)
+
+
+class RMSNorm(nn.Module):
+
+  @nn.compact
+  def __call__(self, x):
+    scale = self.param("scale", nn.initializers.zeros_init(), (x.shape[-1]))
+    var = jnp.mean(jnp.square(x), axis=-1, keepdims=True)
+    normed_inputs = jnp.asarray(x * jnp.reciprocal(jnp.sqrt(var + 1e-06)))
+    normed_inputs = normed_inputs * (1 + scale)
+    return normed_inputs
+
+
+class Embedder(nn.Module):
+  """Embedder module."""
+
+  vocab_size: int
+  embed_dim: int
+
+  def setup(self):
+    self.input_embedding_table = self.param(
+        "input_embedding",
+        nn.initializers.variance_scaling(
+            scale=1.0, mode="fan_in", distribution="normal",
+            in_axis=1, out_axis=0,),
+        (self.vocab_size, self.embed_dim),
+    )
+
+  def encode(self, x):
+    x = self.input_embedding_table[(x,)]
+    x *= jnp.sqrt(self.embed_dim).astype(x.dtype)
+    return x
+
+  def decode(self, x):
+    return jnp.dot(x, self.input_embedding_table.T)
+
+
+class Attention(nn.Module):
+  """Attention module."""
+
+  num_heads: int
+  num_kv_heads: int
+  features: int
+  head_dim: int
+
+  query_pre_attn_norm: str
+  attn_logits_softcap: float | None
+
+  cache_dtype: str | None = None
+
+  def setup(self):
+    if self.num_kv_heads == self.num_heads:
+      self.qkv_einsum = Einsum(
+          shape=(3, self.num_heads, self.features, self.head_dim),
+          w_init=trunc_norm_init(
+              in_axis=(2,), out_axis=(0, 1, 3), batch_axis=()),
+      )
+    else:
+      # MQA / GQA
+      self.q_einsum = Einsum(
+          shape=(self.num_heads, self.features, self.head_dim),
+          w_init=trunc_norm_init(in_axis=(1,), out_axis=(0, 2), batch_axis=()),
+      )
+      self.kv_einsum = Einsum(
+          shape=(2, self.num_kv_heads, self.features, self.head_dim),
+          w_init=trunc_norm_init(
+              in_axis=(2,), out_axis=(0, 1, 3), batch_axis=()),
+      )
+    self.attn_vec_einsum = Einsum(
+        shape=(self.num_heads, self.head_dim, self.features),
+        w_init=trunc_norm_init(in_axis=(0, 1), out_axis=(2,), batch_axis=()),
+    )
+
+  @nn.compact
+  def __call__(self, x, positions, attn_mask, decode, deterministic=True):
+    if self.num_kv_heads == self.num_heads:
+      q, k, v = self.qkv_einsum("BSD,3KDH->3BSKH", x)
+    else:
+      q = self.q_einsum("BTD,NDH->BTNH", x)
+      k, v = self.kv_einsum("BSD,2KDH->2BSKH", x)
+
+    q = _apply_rope(q, positions=positions)
+    if self.query_pre_attn_norm == "rsqrt_head_dim":
+      q *= self.head_dim**-0.5
+    elif self.query_pre_attn_norm == "rsqrt_emb_per_head":
+      q *= (self.features // self.num_heads)**-0.5
+    else:
+      raise ValueError(
+          f"Unknown query_pre_attn_norm: {self.query_pre_attn_norm}"
+      )
+
+    k = _apply_rope(k, positions=positions)
+    if decode:
+      k, v = _update_kv_cache(self, k, v,
+                              cache_size=attn_mask.shape[-1],
+                              cache_dtype=self.cache_dtype)
+
+    q = einops.rearrange(q, "B T (K G) H -> B T K G H", K=self.num_kv_heads)
+    logits = jnp.einsum("BTKGH,BSKH->BKGTS", q, k)
+    logits = logits.astype(jnp.float32)
+
+    if self.attn_logits_softcap:
+      logits = jnp.tanh(logits / self.attn_logits_softcap)
+      logits = logits * self.attn_logits_softcap
+
+    if attn_mask.shape != (q.shape[0], 1, q.shape[1], k.shape[1]):
+      raise ValueError(
+          f"Attention mask with shape {attn_mask.shape} but shapes for q and k "
+          f"are: {q.shape} and {k.shape}"
+      )
+
+    # big_neg = jnp.finfo(logits.dtype).min
+    big_neg = -2.3819763e38  # See gemma/modules.py
+    masked_logits = jnp.where(attn_mask[:, :, None, :, :], logits, big_neg)
+
+    probs = jax.nn.softmax(masked_logits, axis=-1).astype(k.dtype)
+
+    encoded = jnp.einsum("BKGTS,BSKH->BTKGH", probs, v)
+    encoded = einops.rearrange(encoded, "B T K G H -> B T (K G) H")
+    attn_output = self.attn_vec_einsum("BTNH,NHD->BTD", encoded)
+
+    return attn_output
+
+
+class FeedForward(nn.Module):
+  """Feed forward module."""
+
+  features: int
+  hidden_dim: int
+
+  @nn.compact
+  def __call__(self, x):
+    w_gating = self.param(
+        "gating_einsum",
+        trunc_norm_init(in_axis=(1,), out_axis=(0, 2), batch_axis=()),
+        ((2, self.features, self.hidden_dim)),
+    )
+    ff_gate = jnp.dot(x, w_gating[0])
+    gate_value = nn.gelu(ff_gate)
+
+    ff1 = jnp.dot(x, w_gating[1])
+    activations = gate_value * ff1
+
+    w_linear = self.param(
+        "linear",
+        trunc_norm_init(in_axis=(0,), out_axis=(1,), batch_axis=()),
+        (self.hidden_dim, self.features),
+    )
+    outputs = jnp.dot(activations, w_linear)
+
+    return outputs
+
+
+class Block(nn.Module):
+  """Transformer block."""
+
+  num_heads: int
+  num_kv_heads: int
+  embed_dim: int
+  head_dim: int
+  hidden_dim: int
+
+  query_pre_attn_norm: str
+  attn_logits_softcap: float | None
+  post_norms: bool
+
+  dropout: float = 0.0
+  dropout_bdims: tuple[int, ...] = ()
+  cache_dtype: str | None = None
+
+  def setup(self):
+    self.pre_attention_norm = RMSNorm()
+    self.attn = Attention(
+        num_heads=self.num_heads,
+        num_kv_heads=self.num_kv_heads,
+        features=self.embed_dim,
+        head_dim=self.head_dim,
+        cache_dtype=self.cache_dtype,
+        query_pre_attn_norm=self.query_pre_attn_norm,
+        attn_logits_softcap=self.attn_logits_softcap,
+    )
+    self.pre_ffw_norm = RMSNorm()
+    self.mlp = FeedForward(features=self.embed_dim, hidden_dim=self.hidden_dim)
+    if self.dropout:
+      self.drop = nn.Dropout(self.dropout, self.dropout_bdims)
+    else:
+      self.drop = lambda x, _: x
+    if self.post_norms:
+      self.post_attention_norm = RMSNorm()
+      self.post_ffw_norm = RMSNorm()
+
+  def __call__(self, x, unused_scan_arg, positions, attn_mask,
+               decode, deterministic=True):
+    x = nn.with_logical_constraint(x, ("act_batch", "act_len", "act_emb"))
+    inputs_normalized = self.pre_attention_norm(x)
+    attn_output = self.attn(inputs_normalized, positions, attn_mask,
+                            decode, deterministic)
+    if self.post_norms:
+      attn_output = self.post_attention_norm(attn_output)
+    attn_output = self.drop(attn_output, deterministic)
+    attn_output += x
+    residual = attn_output
+    attn_output = self.pre_ffw_norm(attn_output)
+    outputs = self.mlp(attn_output)
+    outputs = self.drop(outputs, deterministic)
+    if self.post_norms:
+      outputs = self.post_ffw_norm(outputs)
+    outputs = residual + outputs
+    return outputs, unused_scan_arg
+
+
+class Model(nn.Module):
+  """gemma model."""
+
+  variant: str
+
+  width: int
+  depth: int
+  mlp_dim: int
+  num_heads: int
+  num_kv_heads: int
+  head_dim: int
+  norm_eps: float
+  vocab_size: int
+
+  query_pre_attn_norm: str = "rsqrt_head_dim"
+  final_logits_softcap: float = 0.0
+  attn_logits_softcap: float = 0.0
+  post_norms: bool = False
+
+  dropout: float = 0.0
+  dropout_bdims: tuple[int, ...] = ()  # Every float is dropped independently.
+  cache_dtype: str | None = None
+
+  # TODO: Wire this in all places needed so that the model can be
+  # run with different activation dtype. For now only float32 runs.
+  embed_dtype: str = "float32"
+
+  scan: bool = False
+  remat_policy: str = "none"
+
+  @nn.compact
+  def __call__(
+      self, tokens, *,
+      embedded_prefix=None,
+      embed_only=False,
+      pre_logits=None,
+      positions=None, mask=None,
+      decode=False, deterministic=True,
+  ):
+    """Embed only, or complete forward pass.
+
+    Args:
+      tokens: Embedded, then and appended to `embedded_prefix`. Can be None.
+      embedded_prefix: Optional prefix that is already embedded.
+      embed_only: Whether to compute embeddings only.
+      pre_logits: If present computes logits from pre_logits and returns.
+      positions: Optional `[B, T]` allows to specify the absolute position of
+        the tokens.
+      mask: Optional attention mask `[B, T, S]`.
+      decode: Whether to use kv-cache. Caller must pass masks and positions.
+      deterministic: Forwarded to all dropout layers.
+
+    Returns:
+      If `embed_only=False`, then `(logits, out)` will be returned.
+      If `embed_only=True`, then the embeddings will be returned.
+    """
+    out = {}
+
+    embedder = Embedder(
+        vocab_size=self.vocab_size,
+        embed_dim=self.width,
+        name="embedder")
+
+    if pre_logits is not None:
+      x = out["pre_logits"] = pre_logits
+      logits = out["logits"] = embedder.decode(x)
+      return logits, out
+
+    x = []
+    if embedded_prefix is not None:
+      x.append(embedded_prefix)
+    if tokens is not None:
+      x.append(embedder.encode(tokens))
+
+    x = jnp.concatenate(x, axis=-2)
+    x = x.astype(self.embed_dtype)
+    batch_size, seq_len, width = x.shape
+
+    if embed_only:
+      return x
+
+    if decode:
+      assert positions is not None and mask is not None, (
+          "Must explicitly pass positions and mask for decoding.")
+
+    if positions is None:
+      positions = jnp.arange(seq_len).astype(jnp.int32)[None, :]
+    assert positions.shape[1] == x.shape[1], (positions.shape, x.shape)
+
+    if mask is None:
+      mask = nn.attention.make_causal_mask(jnp.ones([batch_size, seq_len]))
+    if mask.ndim == 3:
+      mask = mask[:, None, :, :]
+    cache_size = max(seq_len, mask.shape[-1])
+    assert mask.shape == (batch_size, 1, seq_len, cache_size), mask.shape
+
+    if self.remat_policy == "none":
+      block_cls = Block
+    else:
+      block_cls = nn.remat(
+          Block,
+          prevent_cse=not self.scan,
+          static_argnums=(5, 6),  # 0=self, 5=decode, 6=deterministic
+          policy=getattr(jax.checkpoint_policies, self.remat_policy),
+      )
+
+    block_kw = dict(
+        num_heads=self.num_heads,
+        head_dim=self.head_dim,
+        num_kv_heads=self.num_kv_heads,
+        embed_dim=width,
+        hidden_dim=self.mlp_dim,
+        dropout=self.dropout,
+        dropout_bdims=self.dropout_bdims,
+        cache_dtype=self.cache_dtype,
+        query_pre_attn_norm=self.query_pre_attn_norm,
+        attn_logits_softcap=self.attn_logits_softcap,
+        post_norms=self.post_norms,
+    )
+    layers = self.scope.push("layers")  # pytype: disable=attribute-error
+    if self.scan:
+      blocks = [nn.scan(
+          block_cls,
+          # cache has axis 1 since we want leading dimension to be batch size.
+          variable_axes={"params": 0, "cache": 1},
+          split_rngs={"params": True, "dropout": True},
+          in_axes=nn.broadcast,
+          length=self.depth,
+      )(
+          parent=layers, **block_kw
+      )]
+    else:
+      blocks = [
+          block_cls(
+              parent=layers.push(str(layer)),
+              **block_kw,
+          )
+          for layer in range(self.depth)
+      ]
+    unused_scan_arg = ()
+    for block in blocks:
+      x, unused_scan_arg = block(
+          x, unused_scan_arg, positions, mask, decode, deterministic)
+
+    assert x.dtype == jnp.dtype(self.embed_dtype)  # Sanity check.
+    out["encoded"] = x
+
+    x = RMSNorm(name="final_norm")(x)
+    out["pre_logits"] = x
+
+    x = embedder.decode(x)
+    out["logits_pre_norm"] = x
+    if self.final_logits_softcap:
+      x = jnp.tanh(x / self.final_logits_softcap) * self.final_logits_softcap
+    out["logits"] = x
+
+    return x, out
+
+
+_ORBAX_INITS = {}
+_BV_INITS = {}
+
+
+def _load_orbax(path):
+  """Loads and coverts Orbax gemma checkpoint."""
+  checkpointer = orbax.checkpoint.PyTreeCheckpointer()
+  params = checkpointer.restore(path)
+  params = flax.traverse_util.unflatten_dict(params, sep="/")["transformer"]
+  n = sum(1 for k in params if k.startswith("layer_"))
+  params["layers"] = jax.tree.map(
+      lambda *xs: np.stack(xs), *(params.pop(f"layer_{i}") for i in range(n))
+  )
+  mlp = params["layers"]["mlp"]
+  mlp["gating_einsum"] = mlp["gating_einsum"].pop("w")
+  mlp["linear"] = mlp["linear"].pop("w")
+  return params
+
+
+def _del_pad_rows(params):
+  """Some checkpoints have 128 unused padding tokens."""
+  emb = params["embedder"]["input_embedding"]
+  if emb.shape[0] == 256_128:
+    params["embedder"]["input_embedding"] = jax.device_get(emb)[:256_000]
+  assert params["embedder"]["input_embedding"].shape[0] == 256_000
+
+
+def _maybe_transpose_gating_einsum(params):
+  """The `transpose_gating_einsum` case in gemma/modules.py."""
+  mlp = params["layers"]["mlp"]
+  *_, d1, d2 = mlp["gating_einsum"].shape
+  if d1 > d2:
+    *ns, n1, n2 = range(len(mlp["gating_einsum"].shape))
+    mlp["gating_einsum"] = mlp["gating_einsum"].transpose(*ns, n2, n1)
+
+
+def _load_like_bv(params):
+  params = jax.tree.map(lambda x: x, params)
+  _del_pad_rows(params)
+  _maybe_transpose_gating_einsum(params)
+  return params
+
+
+def load(init_params, init_file, model_cfg=None, dont_load=()):
+  """Loads existing weights."""
+  model_cfg = model_cfg or {}
+  variant = model_cfg.get("variant", "gemma_2b")
+  init_variant = f"{init_file} {variant}"
+  if init_variant in _ORBAX_INITS:
+    params = _load_like_bv(_load_orbax(_ORBAX_INITS[init_variant]))
+  elif init_variant in _BV_INITS:
+    params = _load_like_bv(u.load_params(_BV_INITS[init_variant]))
+  else:
+    params = u.load_params(init_file)
+
+  def extend_rows(emb1, target_rows):
+    if (missing_rows := target_rows - emb1.shape[0]) == 0:
+      return emb1
+    assert missing_rows > 0, "You're asking to shrink vocab?!"
+    new_rows = np.random.randn(missing_rows, emb1.shape[1])
+    new_rows = (new_rows * 0.02).astype(emb1.dtype)
+    return np.r_[np.asarray(emb1), new_rows]
+
+  if "vocab_size" in model_cfg:
+    params["embedder"]["input_embedding"] = extend_rows(
+        params["embedder"]["input_embedding"],
+        model_cfg["vocab_size"],
+    )
+
+  return common.merge_params(params, init_params, dont_load)

Tipsomaly/model/big_vision/models/proj/cappa/cappa.py

@@ -0,0 +1,428 @@
+# Copyright 2023 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Model definitions for CapPa (https://arxiv.org/abs/2306.07915).
+
+Used abbreviations for dimension annotations:
+  B: batch size.
+  H: image height.
+  W: image width.
+  P: number of patches (PH/PW: number of patches in height/width dimensions).
+  E: embedding size.
+  L: sequence length of text tokens.
+  V: vocab size.
+"""
+
+from collections.abc import Sequence
+
+from big_vision import utils
+from big_vision.models import common
+from big_vision.models import vit
+import flax
+import flax.linen as nn
+from flax.linen import partitioning
+import jax
+import jax.numpy as jnp
+
+
+def shift_right(x, axis=1, constant_values=0):
+  """Shift to the right on given axis with padding value 0."""
+  pad_widths = [(0, 0)] * len(x.shape)
+  pad_widths[axis] = (1, 0)
+  padded = jnp.pad(x, pad_widths, constant_values=constant_values)
+  # Cuts off the rightmost slice of size along the `axis` dimension.
+  # Note that `list[:-1]`` is the same as `list[slice(-1)]`.
+  return padded[tuple(slice(-1 if i == axis else None) for i in range(x.ndim))]
+
+
+class MlpBlock(nn.Module):
+  """Transformer MLP / feed-forward block with option to deactivate bias."""
+  mlp_dim: int | None = None  # Defaults to 4x input dim
+  dropout: float = 0.0
+  use_bias: bool = True
+
+  @nn.compact
+  def __call__(self, x, deterministic=True):
+    """Applies Transformer MlpBlock module."""
+    inits = dict(
+        kernel_init=nn.initializers.xavier_uniform(),
+        bias_init=nn.initializers.normal(stddev=1e-6),
+    )
+
+    n, l, d = x.shape  # pylint: disable=unused-variable
+    x = nn.Dense(self.mlp_dim or 4 * d, use_bias=self.use_bias, **inits)(x)
+    x = nn.gelu(x)
+    x = nn.Dropout(rate=self.dropout)(x, deterministic)
+    x = nn.Dense(d, use_bias=self.use_bias, **inits)(x)
+    return x
+
+
+class EncoderDecoderBlock(nn.Module):
+  """Transformer encoder-decoder layer."""
+  mlp_dim: int
+  num_heads: int
+  dropout_rate: float = 0.
+  decode: bool = False
+  use_bias: bool = True
+
+  @nn.compact
+  def __call__(self, targets, encoded, decoder_mask=None, deterministic=True):
+    """Applies EncoderDecoder1DBlock module.
+
+    Args:
+      targets: target text embeddings [B, L, E].
+      encoded: encoded image patches from encoder [B, P, E].
+      decoder_mask: decoder self-attention mask.
+      deterministic: bool, deterministic or not (to apply dropout).
+
+    Returns:
+      output after transformer encoder-decoder block [B, L, E].
+    """
+    def wlc(f):
+      dim_names = ("act_batch", "act_len", "act_emb")
+      return nn.with_logical_constraint(f, dim_names)
+
+    # Decoder block.
+    x = wlc(nn.LayerNorm(name="LayerNorm1", use_bias=self.use_bias)(targets))
+    x = wlc(nn.SelfAttention(
+        num_heads=self.num_heads, use_bias=False, broadcast_dropout=False,
+        dropout_rate=self.dropout_rate, decode=self.decode, name="SelfAttn")(
+            x, decoder_mask, deterministic=deterministic))
+    x = wlc(nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic))
+    x = wlc(x + targets)
+
+    if encoded is not None:
+      # Encoder-Decoder block.
+      y = wlc(nn.LayerNorm(name="LayerNorm2", use_bias=self.use_bias)(x))
+      y = wlc(nn.MultiHeadDotProductAttention(
+          num_heads=self.num_heads, use_bias=False, broadcast_dropout=False,
+          dropout_rate=self.dropout_rate, name="CrossAttn")(
+              y, encoded, deterministic=deterministic))
+      y = wlc(
+          nn.Dropout(rate=self.dropout_rate)(y, deterministic=deterministic))
+      y = wlc(y + x)
+    else:
+      y = x
+
+    # MLP block.
+    z = wlc(nn.LayerNorm(name="LayerNorm3", use_bias=self.use_bias)(y))
+    z = wlc(MlpBlock(
+        mlp_dim=self.mlp_dim, dropout=self.dropout_rate, use_bias=self.use_bias,
+        name="MLP")(z, deterministic=deterministic))
+
+    return wlc(y + z), None
+
+
+class Decoder(nn.Module):
+  """Transformer decoder with parallel prediction."""
+  emb_dim: int
+  mlp_dim: int
+  num_heads: int
+  num_layers: int
+  dropout_rate: float = 0.
+  output_vocab_size: int = 32_000
+
+  # Masked prediction training mode
+  masked_pred_prob: float = 0.
+  masking_ratio: float = 0.
+
+  # Whether to use bias in MLP blocks and LN
+  use_bias: bool = True
+
+  scan: bool = False
+  remat_policy: str = "nothing_saveable"
+
+  @nn.compact
+  def __call__(self,
+               encoded,
+               targets,
+               pos_emb,
+               decoder_mask=None,
+               decode=False,
+               deterministic=True,
+               max_decode_length=None):
+    """Applies Transformer model on the inputs.
+
+    Args:
+      encoded: encoded image patches from encoder [B, P, E].
+      targets: target text tokens [B, L].
+      pos_emb: positional embeddings.
+      decoder_mask: decoder self-attention mask.
+      decode: bool, whether to perform fast autoregressive decoding with cache.
+      deterministic: bool, deterministic or not (to apply dropout).
+      max_decode_length: optional max length for positional embeddings.
+
+    Returns:
+      output of a transformer decoder [B, L, V].
+    """
+    y = targets.astype("int32")
+    if not decode:
+      if self.masked_pred_prob > 0.0 and not deterministic:
+        # Binary random variable indicating whether to do masked prediction
+
+        def _add_random_masks(a):
+          # Generate random mask
+          n_masked = int(self.masking_ratio * a.shape[1])
+          mask_locations = jnp.zeros(a.shape[:2], dtype=jnp.int32)
+          mask_locations = mask_locations.at[:, :n_masked].set(1)
+          mask_locations = jax.random.permutation(
+              self.make_rng("dropout"), mask_locations, axis=1, independent=True
+          )
+          # Replace mask locations with mask token index (=vocab_size)
+          a_masked = jnp.where(mask_locations, self.output_vocab_size, a)
+          return a_masked
+
+        def where(mask, x, y):
+          mask = mask.reshape((-1,) + (1,) * (x.ndim - 1))
+          return jnp.where(mask, x, y)
+
+        do_masked_pred = (
+            jax.random.uniform(self.make_rng("dropout"), (len(y),))
+            < self.masked_pred_prob
+        )
+        y = where(do_masked_pred, _add_random_masks(y), shift_right(y))
+        decoder_mask = where(
+            do_masked_pred, jnp.ones_like(decoder_mask), decoder_mask
+        )
+
+      else:
+        y = shift_right(y)
+
+    embed = nn.Embed(
+        self.output_vocab_size + (1 if self.masked_pred_prob > 0.0 else 0),
+        self.emb_dim,
+        name="EmbedTargets",
+        embedding_init=nn.initializers.normal(stddev=1.0),
+    )
+    y = embed(y)
+
+    y = common.AddPositionEmbs(
+        decode=decode, name="PosEmbedTargets")(y, pos_emb)
+    # NOTE: One could apply dropout on the decoder's inputs here. Whether to do
+    # it or not, and if so, what is the best/common way, is to be determined.
+    # y = nn.Dropout(rate=self.dropout_rate)(y, deterministic=deterministic)
+
+    if self.scan:
+      # Mostly followed
+      # https://github.com/google/maxtext/blob/4d99e30b3e0e0cb1d1aa11c7db7fffe18e301498/MaxText/layers.py#L1126
+      # for the scanned version.
+      # 1. remat
+      enc_dec_block_remat = nn.remat(
+          EncoderDecoderBlock,
+          prevent_cse=False,
+          static_argnums=(-1,),
+          policy=getattr(jax.checkpoint_policies, self.remat_policy, None))
+      # 2. scan
+      initializing = self.is_mutable_collection("params")
+      param_scan_axis = 1
+      params_spec = (param_scan_axis if initializing
+                     else partitioning.ScanIn(param_scan_axis))
+      dec_scanned = nn.scan(enc_dec_block_remat,
+                            variable_axes={
+                                "params": params_spec,
+                                "cache": 0,
+                            },
+                            split_rngs={"params": True, "dropout": True},
+                            in_axes=nn.broadcast,
+                            length=self.num_layers)
+      # 3. fprop
+      y, _ = dec_scanned(num_heads=self.num_heads, mlp_dim=self.mlp_dim,
+                         dropout_rate=self.dropout_rate, decode=decode,
+                         use_bias=self.use_bias, name="EncDecBlock")(
+                             y, encoded, decoder_mask, deterministic)
+    else:
+      for lyr in range(self.num_layers):
+        y, _ = EncoderDecoderBlock(
+            num_heads=self.num_heads, mlp_dim=self.mlp_dim,
+            dropout_rate=self.dropout_rate, decode=decode,
+            use_bias=self.use_bias, name=f"EncDecBlock{lyr}")(
+                y, encoded, decoder_mask=decoder_mask,
+                deterministic=deterministic)
+
+    y = nn.LayerNorm(name="LayerNorm")(y)
+
+    logits = nn.Dense(
+        self.output_vocab_size,
+        kernel_init=nn.initializers.zeros,
+        name="LogitsDense",
+    )(y)
+    return logits
+
+
+class Model(nn.Module):
+  """Transformer Model for sequence to sequence translation."""
+  # Encoder/decoder:
+  num_heads: int = 8
+  num_layers: int = 6
+  mlp_dim: int = 2048
+  emb_dim: int = 512
+  enc_dropout_rate: float = 0.
+  vocab_size: int = 32_000
+  seq_len: int = 256
+
+  # Encoder:
+  patches: Sequence[int] = (16, 16)
+  input_seq_len: int = 768
+  posemb_type: str = "learn"
+  patch_dropout: float = 0.
+
+  # Decoder:
+  decoder_num_heads: int = 0
+  decoder_num_layers: int = 0
+  decoder_mlp_dim: int = 0
+  decoder_emb_dim: int = 0
+  dec_dropout_rate: float = 0.
+  # Probability of masked prediction rather than autoregressive prediciton.
+  masked_pred_prob: float = 0.
+  # Masking ratio for masked prediction.
+  masking_ratio: float = 0.
+  # Whether to use bias in decoder MLP blocks and LN.
+  decoder_bias: bool = True
+
+  scan: bool = False
+  remat_policy: str = "nothing_saveable"
+
+  def setup(self):
+
+    self.encoder = vit.Model(
+        patch_size=self.patches,
+        width=self.emb_dim,
+        depth=self.num_layers,
+        mlp_dim=self.mlp_dim,
+        num_heads=self.num_heads,
+        dropout=self.enc_dropout_rate,
+        posemb=self.posemb_type,
+        scan=self.scan,
+        remat_policy=self.remat_policy,
+    )
+
+    self.pos_emb_for_decoder = vit.get_posemb(
+        self,
+        self.posemb_type,
+        (1, self.seq_len),
+        self.decoder_emb_dim or self.emb_dim,
+        "pos_embedding_decoder",
+    )
+    self.decoder = Decoder(
+        num_layers=self.decoder_num_layers or self.num_layers,
+        mlp_dim=self.decoder_mlp_dim or self.mlp_dim,
+        num_heads=self.decoder_num_heads or self.num_heads,
+        dropout_rate=self.dec_dropout_rate,
+        emb_dim=self.decoder_emb_dim or self.emb_dim,
+        output_vocab_size=self.vocab_size,
+        masked_pred_prob=self.masked_pred_prob,
+        masking_ratio=self.masking_ratio,
+        use_bias=self.decoder_bias,
+        scan=self.scan,
+        remat_policy=self.remat_policy,
+    )
+
+  def encode(self, image, train=False, return_enc_features=False):
+    """Encodes input image or embeddings."""
+
+    _, out = self.encoder(image, train=train)
+    encoded = out["encoded"]
+
+    # Return intermediate features if required
+    if return_enc_features:
+      return encoded, out
+
+    return encoded
+
+  def decode(self, encoded, targets, decode=False, train=False,
+             max_decode_length=None):
+    """Applies Transformer decoder-branch on encoded-input and target.
+
+    Args:
+      encoded: encoded image patches from encoder [B, P, E].
+      targets: target text tokens [B, L].
+      decode: whether to prepare and use an autoregressive cache.
+      train: whether it is training.
+      max_decode_length: optional max length for positional embeddings.
+
+    Returns:
+      logits array from transformer decoder [B, L, V].
+    """
+    decoder_mask = None if decode else nn.make_causal_mask(targets)
+    logits = self.decoder(
+        encoded,
+        targets,
+        pos_emb=self.pos_emb_for_decoder,
+        decoder_mask=decoder_mask,
+        decode=decode,
+        deterministic=not train,
+        max_decode_length=max_decode_length)
+    return logits
+
+  def __call__(self, image, text, *, decode=False,
+               train=False, return_enc_features=False):
+    """Applies Transformer model on the inputs.
+
+    Args:
+      image: batch of images [B, H, W, 3].
+      text: batch of tokenized texts [B, L].
+      decode: whether to prepare and use an autoregressive cache.
+      train: whether it is training.
+      return_enc_features: whether to return the encoder features.
+
+    Returns:
+      logits array from full transformer [B, L, V].
+    """
+    if return_enc_features:
+      encoded, out = self.encode(image, train=train, return_enc_features=True)
+      return encoded, out
+
+    encoded = self.encode(image, train=train)
+
+    decoded = self.decode(encoded, text, decode=decode, train=train)
+    return decoded
+
+
+def load(init_params, init_files, model_params=None,
+         dont_load=("head/kernel", "head/bias", "cls")):
+  """Loads params from init checkpoint and merges into init_params."""
+
+  if isinstance(init_files, str):
+    # A shortcut for a single file checkpoint of a vtt model.
+    ckpt_params = utils.load_params(init_files)
+    ckpt_params = flax.training.checkpoints.convert_pre_linen(ckpt_params)
+    ckpt_params = common.merge_params(ckpt_params, init_params, dont_load)
+
+    # Detect attempts to load non-scan checkpoint into scan model if possible.
+    if (model_params.get("scan") and
+        "encoderblock" not in ckpt_params["encoder"]["Transformer"]):
+      raise NotImplementedError("Loading a non-scan checkpoint into a "
+                                "scan model is not supported yet!")
+    if (not model_params.get("scan")
+        and "encoderblock" in ckpt_params["encoder"]["Transformer"]):
+      assert "decoder.*" in dont_load or "decoder/.*" in dont_load, (
+          "Converting scan decoder to a non-scan one is not supported yet!")
+      ckpt_params["encoder"] = utils.jit_cpu()(
+          vit.scan_to_pyloop)(ckpt_params["encoder"])
+
+  else:
+    assert set(init_files) == {"encoder"}, "Only encoder init supported"
+    enc_init = init_files["encoder"]
+    ckpt_params = flax.core.freeze(init_params).unfreeze()
+    vit_params = ckpt_params["encoder"]
+    encoder_params = vit.load(
+        vit_params, enc_init, model_cfg={},
+        dont_load=dont_load)
+    ckpt_params["encoder"] = encoder_params
+
+  ckpt_params["encoder"]["pos_embedding"] = vit.resample_posemb(
+      old=ckpt_params["encoder"]["pos_embedding"],
+      new=init_params["encoder"]["pos_embedding"])
+
+  return ckpt_params

Tipsomaly/model/big_vision/models/proj/clippo/one_tower.py

@@ -0,0 +1,96 @@
+# Copyright 2022 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Model definition to train a single ViT model with the contrastive trainer."""
+
+import importlib
+from typing import Optional, Any
+
+from big_vision import utils
+import flax.linen as nn
+import jax.numpy as jnp
+
+ConfigDict = Any
+
+
+class Model(nn.Module):
+  """Single ViT to encode regular images and text images."""
+  image: Optional[ConfigDict] = None
+  image_model: str = "vit"
+  out_dim: int = 768
+  temperature_init: float = 10.0
+
+  @nn.compact
+  def __call__(self, image, text=None, **kw):
+    """Returns (B, C) image and (B, C) text representations, and some extras."""
+    ztxt, zimg = None, None
+    kw = kw or {}
+
+    image_model = importlib.import_module(
+        f"big_vision.models.{self.image_model}"
+    ).Model(**{"num_classes": self.out_dim, **(self.image or {})}, name="img")  # pylint: disable=not-a-mapping
+
+    def _compute_embedding(input_image, prefix):
+      zemb, out_emb = image_model(input_image, **kw)
+      out = {f"{prefix}/{k}": v for k, v in out_emb.items()}
+
+      # Normalize the embeddings.
+      out[f"{prefix}/norm"] = jnp.linalg.norm(zemb, axis=1, keepdims=True)
+      out[f"{prefix}/normalized"] = zemb = zemb / (out[f"{prefix}/norm"] + 1e-8)
+      return zemb, out
+
+    out = {}
+    if image is not None:
+      zimg, out_img = _compute_embedding(image, "img")
+      out.update(out_img)
+
+    if text is not None:
+      ztxt, out_txt = _compute_embedding(text, "txt")
+      out.update(out_txt)
+
+    temp_init = jnp.log(self.temperature_init)
+    t = self.param("t",
+                   lambda key, shape, dtype: temp_init*jnp.ones(shape, dtype),
+                   (1,), jnp.float32)
+    out["t"] = jnp.exp(t)
+    out["t/parameter"] = t
+
+    return zimg, ztxt, out
+
+
+def load(init_params, init_files, model_cfg, img_load_kw={}):  # pylint: disable=dangerous-default-value
+  """Loads the ViT parameters - adapted from proj/image_text/two_towers.py."""
+  if isinstance(init_files, str):
+    # A shortcut for a single file checkpoint of a two_towers model.
+    init_files = {k: f"{init_files}:{k}" for k in ("img", "t")}
+  else:
+    init_files = {**init_files}  # Shallow copy because we'll pop stuff off.
+
+  restored_params = {**init_params}
+
+  img_init = init_files.pop("image", init_files.pop("img", None))
+  if img_init:
+    restored_params["img"] = importlib.import_module(
+        f"big_vision.models.{model_cfg.image_model}"
+    ).load(init_params["img"], img_init, model_cfg.image, **img_load_kw)
+
+  t_init = init_files.pop("temperature", init_files.pop("t", None))
+  if t_init:
+    restored_params["t"] = utils.load_params(None, t_init)
+
+  assert not init_files, (
+      f"There's something unused left in `config.model_init`. You probably got "
+      f"a typo. Here it is: {init_files}")
+
+  return restored_params

Tipsomaly/model/big_vision/models/proj/flaxformer/bert.py

@@ -0,0 +1,94 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""BERT encoder, optionally loading pre-trained checkpoints."""
+
+import dataclasses
+from typing import Optional
+
+from absl import logging
+from big_vision import utils
+from big_vision.models import common
+import flax
+import flax.linen as nn
+import jax.numpy as jnp
+from tensorflow.io import gfile
+
+from flaxformer.architectures.bert import bert
+from flaxformer.architectures.bert import bert_checkpoint_converter
+from flaxformer.architectures.bert import configs
+
+
+class Model(nn.Module):
+  """BERT encoder with linear projection on last layer CLS token."""
+
+  config: str
+  num_classes: Optional[int] = None
+  head_zeroinit: bool = True
+
+  @nn.compact
+  def __call__(self, text, *, train=False):
+    out = {}
+
+    batch_size, max_len = text.shape
+    bert_model = bert.BertEncoder(**dataclasses.asdict({
+        "base": configs.BertBaseConfig(),
+        "large": configs.BertLargeConfig(),
+    }[self.config]))
+    x = out["transformed"] = bert_model(
+        token_ids=text,
+        position_ids=jnp.tile(
+            jnp.arange(0, max_len, dtype=jnp.int32), [batch_size, 1]),
+        segment_ids=jnp.zeros([batch_size, max_len], dtype=jnp.int32),
+        input_mask=text.astype(jnp.bool_).astype(jnp.int32),
+        enable_dropout=train,
+    )
+
+    x = out["pre_logits"] = x[:, 0]  # CLS token
+
+    if self.num_classes:
+      kw = {"kernel_init": nn.initializers.zeros} if self.head_zeroinit else {}
+      x = out["logits"] = nn.Dense(self.num_classes, name="head", **kw)(x)
+
+    return x, out
+
+
+def load(params, path, model_cfg=None, dont_load=()):
+  """Returns `params` with BERT weights replaced from checkpoint at `path`."""
+  del model_cfg
+
+  checkpoint_path = f"{path}/bert_model.ckpt"
+  if gfile.exists(f"{checkpoint_path}.index"):
+    logging.info("Loading original BERT checkpoint from '%s'", checkpoint_path)
+    params = flax.core.FrozenDict(params).unfreeze()  # Recursive copy.
+    max_len = (
+        params["BertEncoder_0"]["embedder"]["embedders_position_ids"]
+        ["embedding"].shape[0])
+    bert_params, pooler_params = (
+        bert_checkpoint_converter.load_params_from_tf_checkpoint(
+            checkpoint_path=f"{path}/bert_model.ckpt"))
+    del pooler_params
+    if isinstance(bert_params, flax.core.FrozenDict):
+      bert_params = bert_params.unfreeze()
+    bert_params["embedder"]["embedders_position_ids"]["embedding"] = (
+        bert_params["embedder"]["embedders_position_ids"]["embedding"][:max_len]
+    )
+    return common.merge_params(
+        {"BertEncoder_0": bert_params}, params, dont_load)
+
+  logging.info(
+      "Could not find original BERT checkpoint path '%s', "
+      "loading big_vision checkpoint '%s'", checkpoint_path, path)
+  restored_params = utils.load_params(path)
+  return common.merge_params(restored_params, params, dont_load)

Tipsomaly/model/big_vision/models/proj/flaxformer/bert_test.py

@@ -0,0 +1,77 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for bert."""
+
+import tempfile
+
+from big_vision import input_pipeline
+from big_vision.models.proj.flaxformer import bert
+from big_vision.models.proj.flaxformer import bert_test_util
+import big_vision.pp.builder as pp_builder
+import big_vision.pp.ops_general  # pylint: disable=unused-import
+import big_vision.pp.proj.flaxformer.bert_ops  # pylint: disable=unused-import
+import flax
+import jax
+import jax.numpy as jnp
+import tensorflow as tf
+
+
+# BERT vocabulary for testing.
+_BERT_VOCAB = [
+    "[PAD]",
+    "[UNK]",
+    "this",
+    "is",
+    "a",
+    "test",
+    "[CLS]",
+    "[SEP]",
+]
+_TOKEN_LEN = 16
+
+
+class BertTest(tf.test.TestCase):
+
+  def test_load_apply(self):
+    inkey = "text"
+    vocab_path = f"{tempfile.mkdtemp()}/vocab.txt"
+    with open(vocab_path, "w") as f:
+      f.write("\n".join(_BERT_VOCAB))
+    ds2, _ = input_pipeline.make_for_inference(
+        tf.data.Dataset.from_tensor_slices(
+            {inkey: tf.ragged.constant([["this is a test"]])}),
+        num_ex_per_process=[1],
+        preprocess_fn=pp_builder.get_preprocess_fn(
+            f"bert_tokenize(inkey='{inkey}', vocab_path='{vocab_path}', "
+            f"max_len={_TOKEN_LEN})"
+            "|keep('labels')"),
+        batch_size=1,
+    )
+    text = jnp.array(next(iter(ds2))["labels"])
+    model = bert.Model(config="base")
+    variables = model.init(jax.random.PRNGKey(0), text)
+    params = bert.load(flax.core.unfreeze(variables)["params"],
+                       bert_test_util.create_base_checkpoint())
+    x, out = model.apply({"params": params}, text)
+    self.assertAllEqual(jax.tree_map(jnp.shape, x), (1, 768))
+    self.assertAllEqual(
+        jax.tree_map(jnp.shape, out), {
+            "transformed": (1, 16, 768),
+            "pre_logits": (1, 768),
+        })
+
+
+if __name__ == "__main__":
+  tf.test.main()

Tipsomaly/model/big_vision/models/proj/flaxformer/bert_test_util.py

@@ -0,0 +1,261 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for fake BERT checkpoint."""
+
+import tempfile
+
+import tensorflow.compat.v1 as tf
+
+# Checkpoint structure was extracted with the following (Colab) snippet:
+#
+# !wget https://storage.googleapis.com/bert_models/2020_02_20/uncased_L-12_H-768_A-12.zip  # pylint: disable=line-too-long
+# !unzip uncased_L-12_H-768_A-12.zip
+#
+# import tensorflow.compat.v1 as tf
+#
+# ckpt_reader = tf.train.load_checkpoint('bert_model.ckpt')
+# tf_params = {
+#     tf_name: ckpt_reader.get_tensor(tf_name)
+#     for tf_name in ckpt_reader.get_variable_to_dtype_map()
+# }
+#
+# 'shapes_dtypes = {\n%s\n}' % '\n'.join(
+#     f'  "{k}": ({v.shape}, "{v.dtype}"),'
+#     for k, v, in tf_params.items()
+# )
+
+# pylint: disable=line-too-long
+_BASE_SHAPES_DTYPES = {
+    "cls/seq_relationship/output_bias": ((2,), "float32"),
+    "cls/predictions/transform/LayerNorm/gamma": ((768,), "float32"),
+    "cls/predictions/transform/LayerNorm/beta": ((768,), "float32"),
+    "bert/pooler/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_5/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_9/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_9/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_3/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_7/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_9/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_7/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_9/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_9/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_9/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_9/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_9/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_8/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_4/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_8/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_8/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_11/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_11/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_8/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_8/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_2/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_8/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_1/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_8/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_3/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_8/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_8/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_8/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_7/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_7/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_8/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_7/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_7/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_9/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_7/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_7/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_6/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_6/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_6/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_0/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_6/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_7/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_4/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_5/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_2/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_5/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_5/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_5/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_9/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_3/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_8/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_5/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_5/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_5/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_5/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_5/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_5/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_5/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_4/output/dense/bias": ((768,), "float32"),
+    "bert/embeddings/token_type_embeddings": ((2, 768), "float32"),
+    "bert/encoder/layer_4/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_4/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_7/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_4/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_9/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_10/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_6/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_4/attention/self/query/bias": ((768,), "float32"),
+    "cls/seq_relationship/output_weights": ((2, 768), "float32"),
+    "bert/encoder/layer_7/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_4/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_4/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_4/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_3/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_1/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_2/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_8/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_4/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_3/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_4/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_3/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_1/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_3/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_10/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_3/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_1/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_0/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_10/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_3/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_3/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_1/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_3/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_1/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_3/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_2/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_6/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_11/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_2/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_2/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_2/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_2/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_6/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_11/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_6/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_11/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_11/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_11/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_10/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_11/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_6/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_6/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_11/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_10/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_4/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_11/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_10/attention/self/query/bias": ((768,), "float32"),
+    "bert/embeddings/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_2/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_11/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_11/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_5/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_3/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_10/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_10/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/embeddings/word_embeddings": ((30522, 768), "float32"),
+    "bert/encoder/layer_9/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_9/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_6/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_10/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_6/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_1/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_5/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_2/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_0/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_3/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "cls/predictions/output_bias": ((30522,), "float32"),
+    "bert/encoder/layer_0/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_6/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_0/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_2/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_10/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_5/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_4/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_0/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_0/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_10/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_7/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_3/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_2/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_8/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/embeddings/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_1/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_10/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_2/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_6/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_2/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_11/attention/self/value/bias": ((768,), "float32"),
+    "bert/encoder/layer_9/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_0/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_10/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_10/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_1/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_8/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_0/intermediate/dense/bias": ((3072,), "float32"),
+    "bert/encoder/layer_1/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_1/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_7/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_2/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_8/attention/output/dense/bias": ((768,), "float32"),
+    "cls/predictions/transform/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_6/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_5/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_0/attention/self/value/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_7/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_7/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_1/output/dense/kernel": ((3072, 768), "float32"),
+    "bert/encoder/layer_11/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_4/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_1/attention/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_9/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_2/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_0/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_10/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_1/attention/self/query/bias": ((768,), "float32"),
+    "bert/encoder/layer_3/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/encoder/layer_6/attention/output/dense/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_1/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_11/output/dense/bias": ((768,), "float32"),
+    "cls/predictions/transform/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_0/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_11/attention/self/query/kernel": ((768, 768), "float32"),
+    "bert/encoder/layer_0/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_0/attention/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_7/attention/output/LayerNorm/gamma": ((768,), "float32"),
+    "bert/encoder/layer_4/attention/self/key/bias": ((768,), "float32"),
+    "bert/encoder/layer_10/attention/self/key/kernel": ((768, 768), "float32"),
+    "bert/embeddings/position_embeddings": ((512, 768), "float32"),
+    "bert/encoder/layer_1/output/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_9/intermediate/dense/kernel": ((768, 3072), "float32"),
+    "bert/encoder/layer_0/output/LayerNorm/beta": ((768,), "float32"),
+    "bert/pooler/dense/bias": ((768,), "float32"),
+    "bert/encoder/layer_0/attention/output/LayerNorm/beta": ((768,), "float32"),
+}
+# pylint: enable=line-too-long
+
+
+def create_base_checkpoint():
+  """Returns path to fake Bert "base" checkpoint directory (zero init)."""
+  directory = tempfile.mkdtemp()
+  path = f"{directory}/bert_model.ckpt"
+  with tf.Session() as sess:
+    for name, (shape, dtype) in _BASE_SHAPES_DTYPES.items():
+      tf.Variable(tf.zeros(shape, dtype), name=name)
+    saver = tf.train.Saver()
+    sess.run(tf.global_variables_initializer())
+    saver.save(sess, path)
+  return directory

Tipsomaly/model/big_vision/models/proj/flexi/vit.py

@@ -0,0 +1,226 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A version of ViT with flexible seqlen ((internal link))."""
+
+from typing import Optional, Sequence
+
+from absl import logging
+from big_vision import utils
+from big_vision.models import common
+from big_vision.models import vit
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import numpy as np
+import tensorflow as tf
+
+
+def resample_patchemb(old, new_hw):
+  """Resample the weights of the patch embedding kernel to target resolution.
+
+  We resample the patch embedding kernel by approximately inverting the effect
+  of patch resizing. Colab with detailed explanation:
+  (internal link)
+  With this resizing, we can for example load a B/8 filter into a B/16 model
+  and, on 2x larger input image, the result will match.
+  See (internal link)
+  Args:
+    old: original parameter to be resized.
+    new_hw: target shape (height, width)-only.
+  Returns:
+    Resized patch embedding kernel.
+  """
+  assert len(old.shape) == 4, "Four dimensions expected"
+  assert len(new_hw) == 2, "New shape should only be hw"
+  if tuple(old.shape[:2]) == tuple(new_hw):
+    return old
+
+  logging.info("FlexiViT: resize embedding %s to %s", old.shape, new_hw)
+
+  def resize(x_np, new_shape):
+    x_tf = tf.constant(x_np)[None, ..., None]
+    # NOTE: we are using tf.image.resize here to match the resize operations in
+    # the data preprocessing pipeline.
+    x_upsampled = tf.image.resize(
+        x_tf, new_shape, method="bilinear")[0, ..., 0].numpy()
+    return x_upsampled
+
+  def get_resize_mat(old_shape, new_shape):
+    mat = []
+    for i in range(np.prod(old_shape)):
+      basis_vec = np.zeros(old_shape)
+      basis_vec[np.unravel_index(i, old_shape)] = 1.
+      mat.append(resize(basis_vec, new_shape).reshape(-1))
+    return np.stack(mat).T
+
+  resize_mat = get_resize_mat(old.shape[:2], new_hw)
+  resize_mat_pinv = np.linalg.pinv(resize_mat.T)
+
+  def resample_kernel(kernel):
+    resampled_kernel = resize_mat_pinv @ kernel.reshape(-1)
+    return resampled_kernel.reshape(new_hw)
+  v_resample_kernel = jax.vmap(jax.vmap(resample_kernel, 2, 2), 3, 3)
+  return v_resample_kernel(old)
+
+
+class Patchify(nn.Module):
+  """As a class just to match param names with original ViT."""
+
+  patch_size: Sequence[int] = (32, 32)
+  width: int = 768
+  seqhw: Optional[int] = None
+
+  @nn.compact
+  def __call__(self, image, seqhw=None):
+    n, h, w, c = image.shape  # pylint: disable=unused-variable
+
+    w_emb = self.param(
+        "kernel", nn.initializers.normal(stddev=1/np.sqrt(self.width)),
+        (*self.patch_size, c, self.width), image.dtype)
+    b_emb = self.param("bias", nn.initializers.zeros, self.width, image.dtype)
+
+    # Compute required patch-size to reach `seqhw` given `image` size.
+    seqhw = seqhw or self.seqhw
+    if seqhw is None and self.is_initializing():
+      patch_size = self.patch_size
+    else:
+      patch_size = tuple(np.array((h, w)) // np.array((seqhw, seqhw)))
+
+    if patch_size != self.patch_size:
+      w_emb = resample_patchemb(old=w_emb, new_hw=patch_size)
+
+    x = jax.lax.conv_general_dilated(
+        image, w_emb, window_strides=patch_size, padding="VALID",
+        dimension_numbers=("NHWC", "HWIO", "NHWC"))
+    return x + b_emb
+
+
+class _Model(nn.Module):
+  """ViT model."""
+
+  num_classes: int
+  patch_size: Sequence[int] = (32, 32)
+  posemb_size: Sequence[int] = (7, 7)
+  width: int = 768
+  depth: int = 12
+  mlp_dim: Optional[int] = None  # Defaults to 4x input dim
+  num_heads: int = 12
+  posemb: str = "learn"  # Can also be "sincos2d"
+  pool_type: str = "gap"  # Can also be "map" or "tok"
+  head_zeroinit: bool = True
+
+  seqhw: Optional[int] = None
+
+  @nn.compact
+  def __call__(self, image, *, seqhw=None, train=False):
+    out = {}
+
+    x = out["stem"] = Patchify(
+        self.patch_size, self.width, self.seqhw, name="embedding")(image, seqhw)
+
+    # == Flattening + posemb
+    n, h, w, c = x.shape
+    x = jnp.reshape(x, [n, h * w, c])
+
+    pos_emb = vit.get_posemb(
+        self, self.posemb, self.posemb_size, c, "pos_embedding", x.dtype)
+    if pos_emb.shape[1] != h * w:
+      pos_emb = jnp.reshape(pos_emb, (1, *self.posemb_size, c))
+      pos_emb = jax.image.resize(pos_emb, (1, h, w, c), "linear")
+      pos_emb = jnp.reshape(pos_emb, (1, h * w, c))
+
+    x = out["with_posemb"] = x + pos_emb
+
+    # == Optional [cls] token
+    if self.pool_type == "tok":
+      cls = self.param("cls", nn.initializers.zeros, (1, 1, c), x.dtype)
+      x = jnp.concatenate([jnp.tile(cls, [n, 1, 1]), x], axis=1)
+
+    # == Encoder
+    n, l, c = x.shape  # pylint: disable=unused-variable
+
+    x, out["encoder"] = vit.Encoder(
+        depth=self.depth,
+        mlp_dim=self.mlp_dim,
+        num_heads=self.num_heads,
+        name="Transformer")(x)
+    encoded = out["encoded"] = x
+
+    if self.pool_type == "map":
+      x = out["head_input"] = vit.MAPHead(
+          num_heads=self.num_heads, mlp_dim=self.mlp_dim)(x)
+    elif self.pool_type == "gap":
+      x = out["head_input"] = jnp.mean(x, axis=1)
+    elif self.pool_type == "tok":
+      x = out["head_input"] = x[:, 0]
+      encoded = encoded[:, 1:]
+    else:
+      raise ValueError(f"Unknown pool type: '{self.pool_type}'")
+
+    x_2d = jnp.reshape(encoded, [n, h, w, -1])
+
+    out["pre_logits_2d"] = x_2d
+    out["pre_logits"] = x
+
+    if self.num_classes:
+      kw = {"kernel_init": nn.initializers.zeros} if self.head_zeroinit else {}
+      head = nn.Dense(self.num_classes, name="head", **kw)
+      x_2d = out["logits_2d"] = head(x_2d)
+      x = out["logits"] = head(x)
+
+    return x, out
+
+
+def Model(num_classes, *, variant=None, **kw):  # pylint: disable=invalid-name
+  """Factory function, because linen really don't like what I'm doing!"""
+  return _Model(num_classes, **{**vit.decode_variant(variant), **kw})
+
+
+def load(init_params, init_file, model_cfg, dont_load=()):  # pylint: disable=invalid-name because we had to CamelCase above.
+  """Load init from checkpoint, both old model and this one. +Hi-res posemb."""
+  init_file = {**vit.VANITY_NAMES, **VANITY_NAMES}.get(init_file, init_file)
+  restored_params = utils.load_params(init_file)
+
+  restored_params = vit.fix_old_checkpoints(restored_params)
+
+  # Potentially resize the position embedings if seqlen differs.
+  restored_params["pos_embedding"] = vit.resample_posemb(
+      old=restored_params["pos_embedding"],
+      new=init_params["pos_embedding"])
+
+  # Potentially resize the patch embedding kernel.
+  old_patchemb = restored_params["embedding"]["kernel"]
+  restored_params["embedding"]["kernel"] = resample_patchemb(
+      old=old_patchemb, new_hw=model_cfg.patch_size)
+
+  # possibly use the random init for some of the params (such as, the head).
+  restored_params = common.merge_params(restored_params, init_params, dont_load)
+
+  return restored_params
+
+
+# Shortcut names for some canonical paper checkpoints:
+VANITY_NAMES = {
+    # pylint: disable=line-too-long
+    "FlexiViT-L i1k": "gs://big_vision/flexivit/flexivit_l_i1k.npz",
+    "FlexiViT-B i1k": "gs://big_vision/flexivit/flexivit_b_i1k.npz",
+    "FlexiViT-S i1k": "gs://big_vision/flexivit/flexivit_s_i1k.npz",
+    "FlexiViT-B i21k 90ep": "gs://big_vision/flexivit/flexivit_b_i21k_90ep.npz",
+    "FlexiViT-B i21k 300ep": "gs://big_vision/flexivit/flexivit_b_i21k_300ep.npz",
+    "FlexiViT-B i21k 1000ep": "gs://big_vision/flexivit/flexivit_b_i21k_1000ep.npz",
+    "ViT-B/16 i21k": "gs://big_vision/flexivit/vit_b16_i21k_300ep.npz",
+    "ViT-B/30 i21k": "gs://big_vision/flexivit/vit_b30_i21k_300ep.npz",
+    # pylint: enable=line-too-long
+}

Tipsomaly/model/big_vision/models/proj/flexi/vit_test.py

@@ -0,0 +1,127 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the FlexiViT model."""
+
+from absl.testing import absltest
+from big_vision.models.proj.flexi import vit
+import jax
+from jax import config
+from jax import numpy as jnp
+import numpy as np
+import tensorflow as tf
+
+config.update("jax_enable_x64", True)
+
+
+class PatchEmbTest(absltest.TestCase):
+
+  def _test_patch_emb_resize(self, old_shape, new_shape, n_patches=100):
+    # This test verifies that if we resize the input image patch and resample
+    # the patch embedding accordingly, the output does not change.
+    # NOTE: if the image contains more than one patch, then the embeddings will
+    # change due to patch interaction during the resizing.
+    patch_shape = old_shape[:-2]
+    resized_patch_shape = new_shape[:-2]
+    patches = np.random.randn(n_patches, *old_shape[:-1])
+    w_emb = jnp.asarray(np.random.randn(*old_shape))
+
+    old_embeddings = jax.lax.conv_general_dilated(
+        patches, w_emb, window_strides=patch_shape, padding="VALID",
+        dimension_numbers=("NHWC", "HWIO", "NHWC"), precision="highest")
+
+    patch_resized = tf.image.resize(
+        tf.constant(patches), resized_patch_shape, method="bilinear").numpy()
+    patch_resized = jnp.asarray(patch_resized).astype(jnp.float64)
+    w_emb_resampled = vit.resample_patchemb(w_emb, resized_patch_shape)
+    self.assertEqual(w_emb_resampled.shape, new_shape)
+
+    new_embeddings = jax.lax.conv_general_dilated(
+        patch_resized, w_emb_resampled, window_strides=resized_patch_shape,
+        padding="VALID", dimension_numbers=("NHWC", "HWIO", "NHWC"),
+        precision="highest")
+
+    self.assertEqual(old_embeddings.shape, new_embeddings.shape)
+    np.testing.assert_allclose(
+        old_embeddings, new_embeddings, rtol=1e-1, atol=1e-4)
+
+  def test_resize_square(self):
+    out_channels = 256
+    patch_sizes = [48, 40, 30, 24, 20, 16, 15, 12, 10, 8, 6, 5]
+    for s in patch_sizes:
+      old_shape = (s, s, 3, out_channels)
+      for t in patch_sizes:
+        new_shape = (t, t, 3, out_channels)
+        if s <= t:
+          self._test_patch_emb_resize(old_shape, new_shape)
+
+  def test_resize_rectangular(self):
+    out_channels = 256
+    old_shape = (8, 10, 3, out_channels)
+    new_shape = (10, 12, 3, out_channels)
+    self._test_patch_emb_resize(old_shape, new_shape)
+
+    old_shape = (8, 6, 3, out_channels)
+    new_shape = (9, 15, 3, out_channels)
+    self._test_patch_emb_resize(old_shape, new_shape)
+
+    old_shape = (8, 6, 3, out_channels)
+    new_shape = (15, 9, 3, out_channels)
+    self._test_patch_emb_resize(old_shape, new_shape)
+
+  def test_input_channels(self):
+    out_channels = 256
+    for c in [1, 3, 10]:
+      old_shape = (8, 10, c, out_channels)
+      new_shape = (10, 12, c, out_channels)
+      self._test_patch_emb_resize(old_shape, new_shape)
+
+  def _test_works(self, old_shape, new_shape):
+    old = jnp.asarray(np.random.randn(*old_shape))
+    resampled = vit.resample_patchemb(old, new_shape[:2])
+    self.assertEqual(resampled.shape, new_shape)
+    self.assertEqual(resampled.dtype, old.dtype)
+
+  def test_downsampling(self):
+    # NOTE: for downsampling we cannot guarantee that the outputs would match
+    # before and after downsampling. So, we simply test that the code runs and
+    # produces an output of the correct shape and type.
+    out_channels = 256
+    for t in [4, 5, 6, 7]:
+      for c in [1, 3, 5]:
+        old_shape = (8, 8, c, out_channels)
+        new_shape = (t, t, c, out_channels)
+        self._test_works(old_shape, new_shape)
+
+  def _test_raises(self, old_shape, new_shape):
+    old = jnp.asarray(np.random.randn(*old_shape))
+    with self.assertRaises(AssertionError):
+      vit.resample_patchemb(old, new_shape)
+
+  def test_raises_incorrect_dims(self):
+    old_shape = (8, 10, 3, 256)
+    new_shape = (10, 12, 1, 256)
+    self._test_raises(old_shape, new_shape)
+
+    old_shape = (8, 10, 1, 256)
+    new_shape = (10, 12, 3, 256)
+    self._test_raises(old_shape, new_shape)
+
+    old_shape = (8, 10, 3, 128)
+    new_shape = (10, 12, 3, 256)
+    self._test_raises(old_shape, new_shape)
+
+
+if __name__ == "__main__":
+  absltest.main()

Tipsomaly/model/big_vision/models/proj/givt/adaptor.py

@@ -0,0 +1,174 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Invertible adaptor based on iRevNet.
+
+Based on the PyTorch version from:
+https://github.com/jhjacobsen/pytorch-i-revnet/blob/master/models/iRevNet.py
+"""
+
+from typing import Any, Optional, Sequence
+
+from big_vision import utils
+from big_vision.models import common
+from big_vision.models.proj.givt import cnn
+import einops
+import flax.core
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+
+
+def _split(x: jax.Array) -> tuple[jax.Array, jax.Array]:
+  n = x.shape[-1] // 2
+  x1 = x[:, :, :, :n]
+  x2 = x[:, :, :, n:]
+  return x1, x2
+
+
+def _merge(x1: jax.Array, x2: jax.Array) -> jax.Array:
+  return jnp.concatenate((x1, x2), axis=-1)
+
+
+class IRevNetBlock(nn.Module):
+  """iRevNet Block."""
+  first: int = False
+  dropout_rate: float = 0.
+  num_channels: int = 2
+  num_channels_bottleneck: Optional[int] = None
+  num_grps_norm: int = 32
+
+  @nn.compact
+  def _fx2(self, x: jax.Array, train: bool = True) -> jax.Array:
+    if not self.first:
+      y = nn.GroupNorm(num_groups=self.num_grps_norm, name="gn_0")(x)
+      y = nn.relu(y)
+    else:
+      y = x
+
+    ks = (3, 3)     # hardcode kernel-size 3 for now
+    y = nn.Conv(self.num_channels_bottleneck or self.num_channels,
+                kernel_size=ks, padding=1, use_bias=False)(y)
+    y = nn.GroupNorm(num_groups=self.num_grps_norm, name="gn_1")(y)
+    y = nn.relu(y)
+
+    y = nn.Conv(self.num_channels_bottleneck or self.num_channels,
+                kernel_size=ks, padding=1, use_bias=False)(y)
+    y = nn.Dropout(rate=self.dropout_rate, deterministic=(not train))(y)
+    y = nn.GroupNorm(num_groups=self.num_grps_norm, name="gn_2")(y)
+    y = nn.relu(y)
+
+    y = nn.Conv(self.num_channels, kernel_size=ks, padding=1, use_bias=False)(y)
+
+    return y
+
+  def forward(
+      self,
+      x: tuple[jax.Array, jax.Array],
+      train: bool = True,
+  ) -> tuple[jax.Array, jax.Array]:
+    """Bijective block forward."""
+    x1, x2 = x[0], x[1]
+    fx2 = self._fx2(x2, train=train)
+    y1 = fx2 + x1
+    return (x2, y1)
+
+  def inverse(self,
+              x: tuple[jax.Array, jax.Array],
+              train: bool = True
+              ) -> tuple[jax.Array, jax.Array]:
+    """Bijective block inverse."""
+    x2, y1 = x[0], x[1]
+    fx2 = -self._fx2(x2, train=train)
+    x1 = fx2 + y1
+    return (x1, x2)
+
+
+class IRevNet(nn.Module):
+  """iRevNet."""
+  num_blocks: int = 4
+  num_channels: int = 4
+  num_channels_bottleneck: Optional[int] = None
+  dropout_rate: float = 0.0
+
+  def setup(self) -> None:
+    num_grps_norm = min(32, self.num_channels // 2)
+    self.modules = [
+        IRevNetBlock(
+            first=(i == 0),
+            num_channels=self.num_channels // 2,
+            num_channels_bottleneck=(
+                self.num_channels_bottleneck or self.num_channels) // 2,
+            num_grps_norm=num_grps_norm,
+            dropout_rate=self.dropout_rate,
+        )
+        for i in range(self.num_blocks)
+    ]
+
+  def forward(self, x: jax.Array, train: bool = True) -> jax.Array:
+    out = _split(x)
+    for m in self.modules:
+      out = m.forward(out, train=train)
+    out_bij = _merge(out[0], out[1])
+    return out_bij
+
+  def inverse(self, out_bij: jax.Array, train: bool = True) -> jax.Array:
+    out = _split(out_bij)
+    for m in reversed(self.modules):
+      out = m.inverse(out, train=train)
+    out = _merge(out[0], out[1])
+    return out
+
+  def __call__(self, x: jax.Array, train: bool = True) -> jax.Array:
+    return self.forward(x, train=train)
+
+
+class Model(IRevNet):
+  """Wrapper for IRevNet to function as an adaptor in our setup."""
+
+  pixel_shuffle_patch_size: tuple[int, int] = (1, 1)
+
+  def forward(self, x: jax.Array, train: bool = True) -> jax.Array:
+    # (b, code_len, ch) --> (b, h, w, ch) --> (b, code_len, ch)
+    # h, w are the spatial dimensions after space-to-depth transformation
+    h, w = cnn.get_h_w_pixelshuffle(x.shape[1], self.pixel_shuffle_patch_size)
+    x = einops.rearrange(x, "b (h w) c -> b h w c", h=h, w=w)
+    x = super().forward(x, train)
+    x = einops.rearrange(x, "b h w c -> b (h w) c")   # (b, codelen, codeword_d)
+
+    return x
+
+  def inverse(self, out_bij: jax.Array, train: bool = True) -> jax.Array:
+    # (b, code_len, ch) --> (b, h, w, ch) --> (b, code_len, ch)
+    h, w = cnn.get_h_w_pixelshuffle(
+        out_bij.shape[1], self.pixel_shuffle_patch_size)
+    out_bij = einops.rearrange(out_bij, "b (h w) c -> b h w c", h=h, w=w)
+    out_bij = super().inverse(out_bij, train)
+    out_bij = einops.rearrange(out_bij, "b h w c -> b (h w) c")
+
+    return out_bij
+
+
+def load(
+    init_params: Any,
+    init_file: str,
+    model_params: Any = None,
+    dont_load: Sequence[str] = (),
+) -> Any:
+  """Loads params from init checkpoint and merges into init_params."""
+  del model_params
+  ckpt_params = flax.core.unfreeze(utils.load_params(init_file))
+  if init_params is not None:
+    ckpt_params = common.merge_params(ckpt_params, init_params, dont_load)
+  return ckpt_params

Tipsomaly/model/big_vision/models/proj/givt/adaptor_test.py

@@ -0,0 +1,50 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the IRevNet adaptor."""
+
+from big_vision.models.proj.givt import adaptor
+import jax
+from jax import random
+import jax.numpy as jnp
+
+from absl.testing import absltest
+
+
+class AdaptorTest(googletest.TestCase):
+
+  def test_inversion(self):
+    num_channels = 8
+    input_shape = (1, 24, 24, num_channels)
+
+    rng = random.PRNGKey(758493)
+    _, inp_rng, init_rng, data_rng = jax.random.split(rng, 4)
+
+    dummy_x = random.normal(inp_rng, shape=input_shape)
+    real_x = jax.random.normal(data_rng, shape=input_shape)
+
+    model = adaptor.IRevNet(
+        num_blocks=4,
+        num_channels=num_channels,
+        dropout_rate=0.0,
+    )
+    params = model.init(init_rng, dummy_x)
+
+    real_y = model.apply(params, real_x, method=model.forward)
+    real_x_ = model.apply(params, real_y, method=model.inverse)
+    self.assertTrue(jnp.allclose(real_x, real_x_, atol=1e-5))
+
+
+if __name__ == "__main__":
+  googletest.main()

Tipsomaly/model/big_vision/models/proj/givt/cnn.py

@@ -0,0 +1,376 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CNN encoder/decoder architecture based on the VQ-GAN and MaskGIT papers.
+
+Adapted from https://github.com/google-research/maskgit/blob/main/maskgit/nets/vqgan_tokenizer.py. # pylint: disable=line-too-long
+"""
+
+import dataclasses
+import functools
+import math
+from typing import Any, Sequence
+
+from big_vision import utils
+from big_vision.models import common
+from big_vision.models.proj.givt import vae
+
+import einops
+import flax.linen as nn
+import flax.training.checkpoints
+
+import jax
+import jax.numpy as jnp
+
+
+def _get_norm_layer(train, dtype, norm_type="BN"):
+  """Create normalization layers.
+
+  Args:
+    train: Whether to use the layer in training or inference mode.
+    dtype: Layer output type.
+    norm_type: Which normalization to use "BN", "LN", or "GN".
+
+  Returns:
+    An instance of the the layer.
+  """
+  if norm_type == "BN":
+    return functools.partial(
+        nn.BatchNorm,
+        use_running_average=not train,
+        momentum=0.9,
+        epsilon=1e-5,
+        axis_name=None,
+        axis_index_groups=None,
+        dtype=jnp.float32,
+        use_fast_variance=False)
+  elif norm_type == "LN":
+    return functools.partial(nn.LayerNorm, dtype=dtype, use_fast_variance=False)
+  elif norm_type == "GN":
+    return functools.partial(nn.GroupNorm, dtype=dtype, use_fast_variance=False)
+  else:
+    raise NotImplementedError
+
+
+def _tensorflow_style_avg_pooling(x, window_shape, strides, padding: str):
+  """Avg pooling as done by TF (Flax layer gives different results).
+
+  To be specific, Flax includes padding cells when taking the average,
+  while TF does not.
+
+  Args:
+    x: Input tensor
+    window_shape: Shape of pooling window; if 1-dim tuple is just 1d pooling, if
+      2-dim tuple one gets 2d pooling.
+    strides: Must have the same dimension as the window_shape.
+    padding: Either 'SAME' or 'VALID' to indicate pooling method.
+
+  Returns:
+    pooled: Tensor after applying pooling.
+  """
+  pool_sum = jax.lax.reduce_window(x, 0.0, jax.lax.add,
+                                   (1,) + window_shape + (1,),
+                                   (1,) + strides + (1,), padding)
+  pool_denom = jax.lax.reduce_window(
+      jnp.ones_like(x), 0.0, jax.lax.add, (1,) + window_shape + (1,),
+      (1,) + strides + (1,), padding)
+  return pool_sum / pool_denom
+
+
+def _upsample(x, factor=2, method="nearest"):
+  n, h, w, c = x.shape
+  x = jax.image.resize(x, (n, h * factor, w * factor, c), method=method)
+  return x
+
+
+def _dsample(x):
+  return _tensorflow_style_avg_pooling(
+      x, (2, 2), strides=(2, 2), padding="same")
+
+
+def get_h_w_pixelshuffle(hw, pixel_shuffle_patch_size):
+  # Compute h, w after space-to-depth transformation and before flattening,
+  # assuming the imge before space-to-depth transformation was square.
+  ph, pw = pixel_shuffle_patch_size
+  s = int(math.sqrt(hw * ph * pw))
+  h, w = s // ph, s // pw
+  assert h * w == hw, f"Length {hw} incompatible with pixelshuffle ({ph}, {pw})"
+  return h, w
+
+
+class ResBlock(nn.Module):
+  """Basic Residual Block."""
+  filters: int
+  norm_fn: Any
+  conv_fn: Any
+  dtype: int = jnp.float32
+  activation_fn: Any = nn.relu
+  use_conv_shortcut: bool = False
+
+  @nn.compact
+  def __call__(self, x: jax.Array) -> jax.Array:
+    input_dim = x.shape[-1]
+    residual = x
+    x = self.norm_fn()(x)
+    x = self.activation_fn(x)
+    x = self.conv_fn(self.filters, kernel_size=(3, 3), use_bias=False)(x)
+    x = self.norm_fn()(x)
+    x = self.activation_fn(x)
+    x = self.conv_fn(self.filters, kernel_size=(3, 3), use_bias=False)(x)
+    if input_dim != self.filters:
+      if self.use_conv_shortcut:
+        residual = self.conv_fn(
+            self.filters, kernel_size=(3, 3), use_bias=False)(
+                x)
+      else:
+        residual = self.conv_fn(
+            self.filters, kernel_size=(1, 1), use_bias=False)(
+                x)
+    return x + residual
+
+
+class Encoder(nn.Module):
+  """Encoder Blocks."""
+
+  filters: int
+  num_res_blocks: int
+  channel_multipliers: list[int]
+  embedding_dim: int
+  conv_downsample: bool = False
+  norm_type: str = "GN"
+  activation_fn_str: str = "swish"
+  dtype: int = jnp.float32
+
+  def setup(self) -> None:
+    if self.activation_fn_str == "relu":
+      self.activation_fn = nn.relu
+    elif self.activation_fn_str == "swish":
+      self.activation_fn = nn.swish
+    else:
+      raise NotImplementedError
+
+  @nn.compact
+  def __call__(self, x: jax.Array, train: bool = False) -> jax.Array:
+    conv_fn = nn.Conv
+    norm_fn = _get_norm_layer(
+        train=train, dtype=self.dtype, norm_type=self.norm_type)
+    block_args = dict(
+        norm_fn=norm_fn,
+        conv_fn=conv_fn,
+        dtype=self.dtype,
+        activation_fn=self.activation_fn,
+        use_conv_shortcut=False,
+    )
+    x = conv_fn(self.filters, kernel_size=(3, 3), use_bias=False)(x)
+    num_blocks = len(self.channel_multipliers)
+    for i in range(num_blocks):
+      filters = self.filters * self.channel_multipliers[i]
+      for _ in range(self.num_res_blocks):
+        x = ResBlock(filters, **block_args)(x)
+      if i < num_blocks - 1:
+        if self.conv_downsample:
+          x = conv_fn(filters, kernel_size=(4, 4), strides=(2, 2))(x)
+        else:
+          x = _dsample(x)
+    for _ in range(self.num_res_blocks):
+      x = ResBlock(filters, **block_args)(x)
+    x = norm_fn()(x)
+    x = self.activation_fn(x)
+    x = conv_fn(self.embedding_dim, kernel_size=(1, 1))(x)
+    return x
+
+
+class Decoder(nn.Module):
+  """Decoder Blocks."""
+
+  filters: int
+  num_res_blocks: int
+  channel_multipliers: list[int]
+  norm_type: str = "GN"
+  activation_fn_str: str = "swish"
+  output_dim: int = 3
+  dtype: Any = jnp.float32
+
+  def setup(self) -> None:
+    if self.activation_fn_str == "relu":
+      self.activation_fn = nn.relu
+    elif self.activation_fn_str == "swish":
+      self.activation_fn = nn.swish
+    else:
+      raise NotImplementedError
+
+  @nn.compact
+  def __call__(self, x: jax.Array, train: bool = False) -> jax.Array:
+    conv_fn = nn.Conv
+    norm_fn = _get_norm_layer(
+        train=train, dtype=self.dtype, norm_type=self.norm_type)
+    block_args = dict(
+        norm_fn=norm_fn,
+        conv_fn=conv_fn,
+        dtype=self.dtype,
+        activation_fn=self.activation_fn,
+        use_conv_shortcut=False,
+    )
+    num_blocks = len(self.channel_multipliers)
+    filters = self.filters * self.channel_multipliers[-1]
+    x = conv_fn(filters, kernel_size=(3, 3), use_bias=True)(x)
+    for _ in range(self.num_res_blocks):
+      x = ResBlock(filters, **block_args)(x)
+    for i in reversed(range(num_blocks)):
+      filters = self.filters * self.channel_multipliers[i]
+      for _ in range(self.num_res_blocks):
+        x = ResBlock(filters, **block_args)(x)
+      if i > 0:
+        x = _upsample(x, 2)
+        x = conv_fn(filters, kernel_size=(3, 3))(x)
+    x = norm_fn()(x)
+    x = self.activation_fn(x)
+    x = conv_fn(self.output_dim, kernel_size=(3, 3))(x)
+    return x
+
+
+class Model(vae.Model):
+  """CNN Model."""
+
+  filters: int = 128
+  num_res_blocks: int = 2
+  channel_multipliers: list[int] = dataclasses.field(default_factory=list)
+  conv_downsample: bool = False
+  activation_fn: str = "swish"
+  norm_type: str = "GN"
+  output_dim: int = 3
+  dtype: Any = jnp.float32
+  # If True, rescale the input [-1, 1] -> [0, 1] and clip logvar to [-30, 20]
+  malib_ckpt: bool = False
+  pixel_shuffle_patch_size: tuple[int, int] = (1, 1)
+
+  def setup(self) -> None:
+    # Encoder and decoder
+    self.encoder = Encoder(
+        filters=self.filters,
+        num_res_blocks=self.num_res_blocks,
+        channel_multipliers=self.channel_multipliers,
+        norm_type=self.norm_type,
+        activation_fn_str=self.activation_fn,
+        embedding_dim=2 * self.codeword_dim,
+        conv_downsample=self.conv_downsample,
+        dtype=self.dtype,
+        name="cnn_encoder",
+    )
+    self.decoder = Decoder(
+        filters=self.filters,
+        num_res_blocks=self.num_res_blocks,
+        channel_multipliers=self.channel_multipliers,
+        norm_type=self.norm_type,
+        activation_fn_str=self.activation_fn,
+        output_dim=self.output_dim,
+        dtype=self.dtype,
+        name="cnn_decoder",
+    )
+
+  def _maybe_rescale_input(self, x):
+    return (x + 1.0) / 2.0 if self.malib_ckpt else x
+
+  def _maybe_rescale_output(self, x):
+    return 2.0 * x - 1.0 if self.malib_ckpt else x
+
+  def _maybe_clip_logvar(self, logvar):
+    return jnp.clip(logvar, -30.0, 20.0) if self.malib_ckpt else logvar
+
+  def encode(
+      self,
+      x: jax.Array,
+      *,
+      train: bool = False,
+  ) -> tuple[jax.Array, jax.Array]:
+    x = self._maybe_rescale_input(x)
+    x = self.encoder(x, train=train)                  # (2, 16, 16, 64)
+    assert x.shape[1] == x.shape[2], f"Square spatial dims. required: {x.shape}"
+    mu, logvar = jnp.split(x, 2, axis=-1)             # (2, 16, 16, 32) x 2
+    logvar = self._maybe_clip_logvar(logvar)
+
+    def _space_to_depth(z):
+      ph, pw = self.pixel_shuffle_patch_size
+      return einops.rearrange(
+          z, "b (h ph) (w pw) c -> b (h w) (c ph pw)",
+          ph=ph, pw=pw
+      )   # (2, 256 // (ph * pw), 64 * ph * pw)
+
+    mu, logvar = _space_to_depth(mu), _space_to_depth(logvar)
+
+    return mu, logvar
+
+  def decode(self, x: jax.Array, train: bool = False) -> jax.Array:
+    # Decode
+    ph, pw = self.pixel_shuffle_patch_size
+    h, w = get_h_w_pixelshuffle(x.shape[1], (ph, pw))
+
+    x = einops.rearrange(
+        x, "b (h w) (c ph pw) -> b (h ph) (w pw) c",
+        h=h, w=w,
+        ph=ph, pw=pw
+    )  # (2, 16, 16, 32)
+    x = self.decoder(x, train=train)   # (2, 256, 256, 3)
+    x = self._maybe_rescale_output(x)
+    x = jnp.clip(x, -1.0, 1.0)
+
+    return x
+
+
+def load(
+    init_params: Any,
+    init_file: str,
+    model_params: Any = None,
+    dont_load: Sequence[str] = (),
+    malib_ckpt: bool = False,
+    use_ema_params: bool = False,
+) -> Any:
+  """Loads params from init checkpoint and merges into init_params.
+
+  Args:
+    init_params: pytree with (previously initialized) model parameters.
+    init_file: Path of the checkpoint to load.
+    model_params: Dict containing the model config.
+    dont_load: Sequence of (flattened) parameter names which should not be
+      loaded.
+    malib_ckpt: Whether the given init_file is a malib checkpoint.
+    use_ema_params: Whether to load the EMA params (for malib checkpoints).
+
+  Returns:
+    pytree containing the loaded model parameters.
+  """
+  # `model_params` is unused here, but we still include it to conform with the
+  # general big_vision interface, cf. the core models in big_vision/models/.
+  del model_params
+
+  assert malib_ckpt or (not use_ema_params), (
+      "Loading EMA parameters is only supported for malib checkpoints.")
+
+  if malib_ckpt:
+    # Locally disable transfer guard since restore_checkpoint does not allow for
+    # fine-grained sharding control.
+    with jax.transfer_guard("allow"):
+      vaegan_params = flax.training.checkpoints.restore_checkpoint(
+          init_file, None)
+    vaegan_params_flat = utils.tree_flatten_with_names(vaegan_params)[0]
+    prefix_old = "ema_params/" if use_ema_params else "g_params/"
+    vaegan_params_flat = [(k.replace(prefix_old, "cnn_"), v)
+                          for k, v in vaegan_params_flat if prefix_old in k]
+    params = utils.tree_unflatten(vaegan_params_flat)
+  else:
+    params = flax.core.unfreeze(utils.load_params(init_file))
+
+  if init_params is not None:
+    params = common.merge_params(params, init_params, dont_load)
+  return params

Tipsomaly/model/big_vision/models/proj/givt/decode.py

@@ -0,0 +1,386 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Autorgregressive sampler for GIVT."""
+
+import functools
+from typing import Any, Optional
+
+from big_vision.models.proj.givt import parallel_decode
+import flax
+from flax import linen as nn
+import jax
+from jax import lax
+from jax import numpy as jnp
+import ml_collections
+
+
+def _sample_gmm(
+    gmm_pdf,
+    *,
+    rng,
+    cfg_inference_weight=None,
+    gmm_pdf_uncond=None,
+):
+  """Draw a single sample from a GMM."""
+  if cfg_inference_weight is not None:
+    assert gmm_pdf_uncond is not None
+    gmm_pdf = parallel_decode.CFGDensity(
+        gmm_pdf, gmm_pdf_uncond, w=cfg_inference_weight, rng=rng
+    )
+  samples = gmm_pdf.sample(seed=rng)
+  logprobs = gmm_pdf.log_prob(samples)
+  if logprobs.ndim == 2:
+    logprobs = logprobs[..., None]
+  return samples, logprobs
+
+
+# Beam search reshaping utils
+def _flatten_samples_dim(x):
+  """Flattens samples dimension into batch dimension."""
+  if x.ndim == 0:  # ignore scalars (e.g. cache index)
+    return x
+  return x.reshape((x.shape[0] * x.shape[1],) + x.shape[2:])
+
+
+def _unflatten_samples_dim(x, batch_size, num_samples):
+  """Unflattens first dimension into batch and samples dimensions."""
+  if x.ndim == 0:  # ignore scalars (e.g. cache index)
+    return x
+  assert batch_size * num_samples == x.shape[0]
+  return x.reshape((batch_size, num_samples) + x.shape[1:])
+
+
+def _cache_map(fn, cache, scan=False):
+  """Maps function over cache."""
+  if scan:
+    # Assuming the chache is scanned over the first dimension, we apply a map
+    # function over this dimension for scanned models
+    fn_mod = lambda x: jax.lax.map(fn, x) if x.ndim > 0 else fn(x)
+  else:
+    fn_mod = fn
+
+  frozen = isinstance(cache, flax.core.FrozenDict)
+  if frozen:
+    cache = flax.core.unfreeze(cache)
+  flat_cache = flax.traverse_util.flatten_dict(cache)
+  # Exclude cached relative position bias from beam expansion, etc.
+  keyvals = {k: v for k, v in flat_cache.items() if k[-1] != "cached_bias"}
+  keyvals = jax.tree_map(fn_mod, keyvals)
+  flat_cache.update(keyvals)
+  new_cache = flax.traverse_util.unflatten_dict(flat_cache)
+  if frozen:
+    new_cache = flax.core.freeze(new_cache)
+  return new_cache
+
+
+@flax.struct.dataclass
+class LoopState:
+  """Internal state of the sampling loop."""
+  # Terminology
+  # b:  batch size
+  # nb: number of beams
+  # nf: number of fans
+  # s:  seaquence length
+  # d:  feature dimension
+  rng: jnp.ndarray        # PRNGKey of the loop state.
+  cache: Any              # Cache for fast auto-regressive decoding.
+  sequences: jnp.ndarray  # (b * nb, s, d)
+  logprobs: jnp.ndarray   # (b * nb, s, d)
+  cache_u: Any            # Uncond cache if cfg, otherwise None
+
+
+def _create_cache(
+    labels,
+    model,
+    init_sequence,
+    params,
+    encoded,
+    uncond=False,
+):
+  """Creates the cache and returns initial logits."""
+  if uncond:
+    assert labels is not None  # Need labels for CFG!
+    drop_labels = jnp.ones((labels.shape[0],), dtype=jnp.bool_)
+  else:
+    drop_labels = None
+
+  def init_cache(model):
+    return model.decode(
+        init_sequence, labels, encoded, decode=True, drop_labels=drop_labels
+    )
+
+  cache = nn.apply(init_cache, model, mutable=True)(params)[1]["cache"]
+
+  def prefill_cache(model):
+    return model.prefill(
+        labels, init_sequence.shape[0], encoded, drop_labels=drop_labels
+    )
+
+  # prefill class label or BOS token
+  prefill_logits, aux = nn.apply(prefill_cache, model, mutable=True)(
+      {"params": params["params"], "cache": cache})
+  cache = aux["cache"]
+  return cache, prefill_logits
+
+
+def generate(
+    params: Any,
+    seed: jax.Array,
+    *,
+    model: nn.Module,
+    seq_len: int,
+    feature_dim: int,
+    labels: Optional[jnp.ndarray] = None,
+    cond_image: Optional[jnp.ndarray] = None,
+    batch_size: Optional[int] = None,
+    config: Optional[ml_collections.ConfigDict] = None,
+) -> tuple[jax.Array, jax.Array]:
+  """Sampling loop for GIVT."""
+  if model.style != "ar":  # pytype: disable=wrong-arg-types
+    raise ValueError(f"Invalid style: {model.style}")
+  if model.has_encoder != (cond_image is not None):
+    raise ValueError("Need cond_image if and only if the model has an encoder!")
+
+  assert labels is not None or batch_size, (
+      "Please provide either labels or batch_size.")
+
+  config = config or {}
+  config = dict(config)  # copy
+
+  # For sampling, we support keep_gt (a bool mask), and gt (ground truth)
+  # tokens to use instead of samples.
+  keep_gt = config.pop("keep_gt", None)
+  gt = config.pop("gt", None)
+
+  if isinstance(seed, int):
+    seed = jax.random.PRNGKey(seed)
+
+  beam_size = config.pop("beam_size", 1)
+  fan_size = config.pop("fan_size", 1)
+
+  if labels is not None:
+    batch_size = labels.shape[0]
+    # fold beams into batch dimension
+    labels = labels.repeat(beam_size, axis=0)
+
+  # initialize sequence and logprobs (we track per feature dim logprobs)
+  init_sequence = jnp.zeros((batch_size * beam_size, seq_len, feature_dim))
+  init_logprobs = jnp.zeros_like(init_sequence)
+
+  if cond_image is not None:
+    # embed conditioning image if provided
+    def encode_cond_img(model, cond_img):
+      return model.encode(cond_img)
+    encoded = nn.apply(encode_cond_img, model)(params, cond_image)
+    encoded = jnp.repeat(encoded, beam_size, axis=0)
+  else:
+    encoded = None
+
+  cache, prefill_logits = _create_cache(
+      labels, model, init_sequence, params, encoded
+  )
+
+  cfg_inference_weight = config.pop("cfg_inference_weight", None)
+  if cfg_inference_weight == 0.0:
+    cfg_inference_weight = None
+  cfg = cfg_inference_weight is not None
+
+  get_pdf = functools.partial(
+      model.get_pdf,
+      temperature_scales=config.pop("temp", None),
+      temperature_probs=config.pop("temp_probs", None),
+  )
+
+  # setup sampling function
+  sample = functools.partial(
+      _sample_gmm, cfg_inference_weight=cfg_inference_weight
+  )
+
+  # draw first output token
+  pdf_first = get_pdf(prefill_logits)
+  rng_first, rng = jax.random.split(seed)
+
+  if cfg:
+    assert beam_size == 1 and fan_size == 1  # CFG + Beam not supported.
+    cache_u, prefill_logits_u = _create_cache(
+        labels, model, init_sequence, params, encoded, uncond=True
+    )
+    pdf_first_u = get_pdf(prefill_logits_u)
+  else:
+    cache_u = None
+    pdf_first_u = None
+
+  tokens_first, logprobs_first = sample(
+      pdf_first, rng=rng_first, gmm_pdf_uncond=pdf_first_u
+  )
+  init_sequence = init_sequence.at[:, 0].set(tokens_first.squeeze(axis=1))
+  init_logprobs = init_logprobs.at[:, 0].set(logprobs_first.squeeze(axis=1))
+
+  def tokens_to_logits(tokens, cache, uncond=False):
+    if uncond:
+      drop_labels = jnp.ones((labels.shape[0],), dtype=jnp.bool_)
+    else:
+      drop_labels = None
+
+    def decode_step(model, tokens):
+      return model.decode(tokens, labels, encoded,
+                          decode=True, drop_labels=drop_labels)
+
+    logits, aux = nn.apply(decode_step, model, mutable=True)(
+        {"params": params["params"], "cache": cache}, tokens)
+    return logits, aux["cache"]
+
+  init_state = LoopState(
+      cache=cache,
+      sequences=init_sequence,  # (b * nb, s, d)
+      logprobs=init_logprobs,   # (b * nb, s, d)
+      rng=rng,
+      cache_u=cache_u,
+  )
+
+  rand_top_k = config.pop("rand_top_k", False)
+  rand_top_k_temp = config.pop("rand_top_k_temp", 1.0)
+
+  assert not config, f"Sampling config is expected to be empty: {config}"
+
+  def sampling_iteration(i, state):
+    rng_sampling, rng_local = jax.random.split(state.rng)
+    cur_tokens = state.sequences[:, i][:, None]
+    # (b * nb, d)
+    cur_logits, cache = tokens_to_logits(cur_tokens, state.cache)
+
+    # (b, nb, d)
+    cur_logits = _unflatten_samples_dim(
+        cur_logits, batch_size, beam_size).squeeze(axis=2)
+
+    # (b, nb * nf, d)
+    cur_pdf = get_pdf(cur_logits.repeat(fan_size, axis=1))
+
+    if cfg:
+      cur_logits_u, cache_u = tokens_to_logits(
+          cur_tokens, state.cache_u, uncond=True
+      )
+      cur_logits_u = _unflatten_samples_dim(
+          cur_logits_u, batch_size, beam_size).squeeze(axis=2)
+      cur_pdf_u = get_pdf(cur_logits_u.repeat(fan_size, axis=1))
+      new_tokens, new_logprobs = sample(
+          cur_pdf, rng=rng_sampling, gmm_pdf_uncond=cur_pdf_u
+      )
+    else:
+      new_tokens, new_logprobs = sample(cur_pdf, rng=rng_sampling)
+      cache_u = None
+
+    if gt is not None:
+      assert keep_gt is not None
+      new_tokens = jnp.where(keep_gt[i], gt[:, i, :][:, None], new_tokens)
+
+    # Skip beam search if not needed
+    if beam_size == fan_size == 1:
+      sampled_tokens = new_tokens.squeeze(axis=1)
+      sequences = state.sequences.at[:, i + 1].set(sampled_tokens)
+      return LoopState(
+          cache=cache,
+          rng=rng_local,
+          sequences=sequences,
+          logprobs=state.logprobs,
+          cache_u=cache_u,
+      )
+
+    # (b, nb, s, d)
+    logprobs = _unflatten_samples_dim(state.logprobs, batch_size, beam_size)
+    cur_logprobs = logprobs[:, :, i]  # (b, nb, d)
+    # (b, nb * nf, d)
+    new_logprobs = new_logprobs + cur_logprobs.repeat(fan_size, axis=1)
+    beam_logprobs = new_logprobs.sum(axis=-1)  # (b, nb * nf)
+
+    if rand_top_k:
+      # randomize top-k sampling via sampling from a categorical distribution
+      def stoc_top_k(r, x, p):
+        return jax.random.choice(r, x, shape=(beam_size,), replace=False, p=p)
+      # construct index grid
+      index_grid = jnp.arange(beam_logprobs.shape[1], dtype=jnp.int32)
+      # (b, nb * nf)
+      index_grid = index_grid[None].repeat(beam_logprobs.shape[0], axis=0)
+      top_k_rng, rng_local = jax.random.split(rng_local)
+      top_k_rng = jax.random.split(top_k_rng, beam_logprobs.shape[0])
+      # vmap categorical sampling
+      top_beam_fan_indices = jax.vmap(stoc_top_k, in_axes=(0, 0, 0))(
+          top_k_rng,
+          index_grid,
+          nn.softmax(beam_logprobs / rand_top_k_temp, axis=-1))
+    else:
+      _, top_beam_fan_indices = lax.top_k(beam_logprobs, k=beam_size)  # (b, nb)
+
+    top_beam_indices = top_beam_fan_indices // fan_size
+
+    def _gather_beams(x):
+      if x.ndim == 0:
+        return x
+      # checkify.check(jnp.all(top_beam_indices < x.shape[1]),
+      #                f"`take_along_axis` out of bounds in `_gather_beams`: "
+      #                f"{top_beam_indices.max()} vs. {x.shape[1]}")
+      # (b, nb, 1 ... 1)
+      expanded_indices = top_beam_indices.reshape(
+          top_beam_indices.shape + (1,) * (x.ndim - 2))
+      return jnp.take_along_axis(x, expanded_indices, axis=1)
+
+    def _gather_tokens(x):
+      # (b, nb * nf, d) -> (b, nb, d)
+      # checkify.check(jnp.all(top_beam_fan_indices < x.shape[1]),
+      #                f"`take_along_axis` out of bounds in `_gather_tokens`: "
+      #                f"{top_beam_fan_indices.max()} vs. {x.shape[1]}")
+      return jnp.take_along_axis(x, top_beam_fan_indices[..., None], axis=1)
+    # (b, nb, s, d)
+    sequences = _unflatten_samples_dim(state.sequences, batch_size, beam_size)
+    sequences = _gather_beams(sequences)  # (b, nb, s, d)
+    sequences = sequences.at[:, :, i + 1].set(_gather_tokens(new_tokens))
+    # (b, nb, s, d)
+    sequences = _flatten_samples_dim(sequences)
+
+    logprobs = _gather_beams(logprobs)
+    logprobs = logprobs.at[:, :, i + 1].set(_gather_tokens(new_logprobs))
+    logprobs = _flatten_samples_dim(logprobs)
+
+    scanned_cache = getattr(model, "scan", False)
+    cache = _cache_map(
+        lambda x: _unflatten_samples_dim(x, batch_size, beam_size),
+        cache, scanned_cache)
+    cache = _cache_map(_gather_beams, cache, scanned_cache)
+    cache = _cache_map(_flatten_samples_dim, cache, scanned_cache)
+
+    if cfg:
+      assert cache_u is not None
+      cache_u = _cache_map(
+          lambda x: _unflatten_samples_dim(x, batch_size, beam_size),
+          cache_u, scanned_cache
+      )
+      cache_u = _cache_map(_gather_beams, cache_u, scanned_cache)
+      cache_u = _cache_map(_flatten_samples_dim, cache_u, scanned_cache)
+    else:
+      assert cache_u is None
+
+    return LoopState(
+        cache=cache,
+        rng=rng_local,
+        sequences=sequences,
+        logprobs=logprobs,
+        cache_u=cache_u,
+    )
+
+  final_state = lax.fori_loop(0, seq_len, sampling_iteration, init_state)
+  final_logprobs = final_state.logprobs[::beam_size][:, -1].sum(axis=-1)
+
+  # return top beams and corresponding log probs
+  return final_state.sequences[::beam_size], final_logprobs

Tipsomaly/model/big_vision/models/proj/givt/decode_test.py

@@ -0,0 +1,121 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from absl.testing import parameterized
+from big_vision.models.proj.givt import decode
+from big_vision.models.proj.givt import givt
+import jax
+import jax.numpy as jnp
+
+from absl.testing import absltest
+
+
+_BATCH_SIZE = 2
+_OUT_DIM = 4
+_IMG_DIM = 8
+_PATCH_SIZE = 2
+_SEQ_LEN = _IMG_DIM // _PATCH_SIZE * _IMG_DIM // _PATCH_SIZE
+_NUM_MIXTURES = 4
+
+
+def _make_test_model(**overwrites):
+  config = dict(
+      num_heads=2,
+      num_decoder_layers=1,
+      mlp_dim=64,
+      emb_dim=16,
+      patches=(_PATCH_SIZE, _PATCH_SIZE),
+      input_size=(_IMG_DIM, _IMG_DIM),
+      seq_len=_SEQ_LEN,
+      out_dim=_OUT_DIM,
+      num_mixtures=_NUM_MIXTURES,
+      style="ar",
+  )
+  config.update(overwrites)
+  return givt.Model(**config)
+
+
+class DecodeTest(parameterized.TestCase):
+
+  def _make_model(self, **overwrites):
+    model = _make_test_model(**overwrites)
+    sequence = jax.random.uniform(
+        jax.random.PRNGKey(0), (_BATCH_SIZE, _SEQ_LEN, _OUT_DIM)
+    )
+    labels = jax.random.uniform(
+        jax.random.PRNGKey(0), (_BATCH_SIZE,), maxval=10
+    ).astype(jnp.int32)
+    variables = model.init(
+        jax.random.PRNGKey(0),
+        sequence,
+        labels,
+        train=False,
+        image=jnp.zeros((_BATCH_SIZE, _IMG_DIM, _IMG_DIM, 3), dtype=jnp.float32)
+        if model.has_encoder
+        else None,
+    )
+    return model, variables
+
+  def _test_model(self, rng, model, variables, config):
+    labels = jnp.ones((_BATCH_SIZE,), dtype=jnp.int32)
+    if model.has_encoder:
+      cond_image = jnp.zeros(
+          (_BATCH_SIZE, _IMG_DIM, _IMG_DIM, 3), dtype=jnp.float32
+      )
+    else:
+      cond_image = None
+    result, logprobs = decode.generate(
+        params=variables,
+        seed=rng,
+        seq_len=_SEQ_LEN,
+        feature_dim=_OUT_DIM,
+        labels=labels,
+        model=model,
+        config=config,
+        cond_image=cond_image,
+    )
+    # TODO: More expressive tests? Eg for causality, and caching.
+    self.assertEqual(result.shape, (_BATCH_SIZE, _SEQ_LEN, _OUT_DIM))
+    self.assertTrue(jnp.allclose(logprobs, jnp.zeros_like(logprobs), atol=1e-5))
+
+  @parameterized.product(
+      rng_seed=[1, 2],
+      encoder=[True, False],
+  )
+  def test_simple(self, rng_seed, encoder):
+    rng = jax.random.PRNGKey(rng_seed)
+    model, variables = self._make_model(
+        num_layers=1 if encoder else 0
+    )
+    assert model.has_encoder == encoder
+    self._test_model(rng, model, variables, config={})
+
+  @parameterized.product(
+      rng_seed=[1, 2],
+      cfg_inference_weight=[0.0, 1.0, 3.0],
+      per_channel_mixtures=[True, False],
+  )
+  def test_cfg(self, rng_seed, cfg_inference_weight, per_channel_mixtures):
+    rng = jax.random.PRNGKey(rng_seed)
+    model, variables = self._make_model(
+        num_mixtures=1 if per_channel_mixtures else 3,
+        drop_labels_probability=0.1,
+        per_channel_mixtures=per_channel_mixtures,
+    )
+    config = {"cfg_inference_weight": cfg_inference_weight}
+    self._test_model(rng, model, variables, config)
+
+
+if __name__ == "__main__":
+  googletest.main()

Tipsomaly/model/big_vision/models/proj/givt/givt.py

@@ -0,0 +1,820 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Decoder-only and encoder-decoder GIVT model.
+
+Used abbreviations for dimension annotations:
+  B: batch size.
+  E: embedding size.
+  L: (soft) token sequence length.
+  D: soft token dimension.
+  P: number of patches (extracted by a ViT encoder in GIVT-based UViM)
+"""
+
+import enum
+import itertools
+from typing import Literal, Optional, Sequence, Any, Mapping
+
+from absl import logging
+from big_vision import utils
+from big_vision.models import common
+from big_vision.models import vit
+import distrax
+import einops
+import flax.linen as nn
+from flax.linen import partitioning
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+class _SpecialLabel(enum.Enum):
+
+  MASK = "mask"
+  NOMASK = "nomask"
+  REPLACE = "replace"
+  NOLABEL = "nolabel"  # For CFG
+
+
+def _random_mask_with_ratios(rng, ratios: jax.Array, seq_len: int):
+  """Generates masks where a fraction of tokens is uncovered.
+
+  Args:
+    rng: RNG.
+    ratios: Ratios, must be a 1D matrix of shape (B,). Values must be in
+      [0, 1], and indicate at ratios[i] how many of the i-th tokens are
+      uncovered (ie. equal to `True`).
+    seq_len: How many tokens this mask has to cover.
+
+  Returns:
+    Mask of dtype bool, shape (B, L).
+
+  Raises:
+    ValueError: Incorrect inputs.
+  """
+  if ratios.ndim != 1:
+    raise ValueError("Ratios must have shape (B,)!")
+  ratios = jnp.clip(ratios, 0, 1)
+  indices = jnp.arange(seq_len, dtype=jnp.float32)  # Shape: (L,)
+  ratios = ratios[:, jnp.newaxis] * seq_len  # Shape: (B, 1)
+  # This is a binary array where the first ratios * seq_len positions are True
+  mask = (indices < ratios).astype(jnp.bool_)  # Shape: (B, L)
+  # Shuffle to a actual mask.
+  return jax.random.shuffle(rng, mask, axis=-1)
+
+
+def apply_mask_schedule(ratio: float | jax.Array, method: str) -> jax.Array:
+  """Generate a mask rate by scheduling mask functions R."""
+  if method == "cosine":
+    mask_ratio = jax.lax.cos(jnp.pi / 2. * ratio)
+  elif "pow:" in method:
+    exponent = float(method.replace("pow:", ""))
+    mask_ratio = 1. - ratio**exponent
+  else:
+    raise NotImplementedError(method)
+  # Clamps mask into [epsilon, 1)
+  mask_ratio = jnp.clip(mask_ratio, 1e-6, 1.)
+  return mask_ratio
+
+
+class EncoderDecoderBlock(nn.Module):
+  """Transformer encoder-decoder layer."""
+  mlp_dim: int
+  num_heads: int
+  dropout_rate: float = 0.
+  decode: bool = False
+
+  @nn.compact
+  def __call__(
+      self,
+      targets: jax.Array,
+      encoded: jax.Array | None = None,
+      decoder_mask: jax.Array | None = None,
+      deterministic: bool = True,
+  ) -> tuple[jax.Array, jax.Array]:
+    """Applies EncoderDecoderBlock module.
+
+    Args:
+      targets: target text embeddings [B, L, D].
+      encoded: encoded image patches from encoder [B, P, E].
+      decoder_mask: decoder self-attention mask.
+      deterministic: bool, deterministic or not (to apply dropout).
+
+    Returns:
+      output after transformer encoder-decoder block [B, L, E].
+    """
+    # Helper function for axis annotation.
+    def wlc(f):
+      dim_names = ("act_batch", "act_len", "act_emb")
+      return nn.with_logical_constraint(f, dim_names)
+    # Decoder block.
+    x = wlc(nn.LayerNorm(name="LayerNorm1", use_bias=False)(targets))
+    x = wlc(nn.SelfAttention(
+        num_heads=self.num_heads, use_bias=False, broadcast_dropout=False,
+        dropout_rate=self.dropout_rate, decode=self.decode, name="SelfAttn")(
+            x, decoder_mask, deterministic=deterministic))
+    x = wlc(nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic))
+    x = wlc(x + targets)
+
+    if encoded is None:
+      y = x
+    else:
+      # Encoder-Decoder block.
+      y = wlc(nn.LayerNorm(name="LayerNorm2", use_bias=False)(x))
+      y = wlc(nn.MultiHeadDotProductAttention(
+          num_heads=self.num_heads, use_bias=False, broadcast_dropout=False,
+          dropout_rate=self.dropout_rate, name="CrossAttn")(
+              y, encoded, deterministic=deterministic))
+      y = wlc(
+          nn.Dropout(rate=self.dropout_rate)(y, deterministic=deterministic))
+      y = wlc(y + x)
+
+    # MLP block.
+    z = wlc(nn.LayerNorm(name="LayerNorm3", use_bias=False)(y))
+    z = wlc(vit.MlpBlock(mlp_dim=self.mlp_dim, dropout=self.dropout_rate,
+                         name="MLP")(z, deterministic=deterministic))
+
+    # nn.scan requires a carry (second element in tuple)
+    out = wlc(y + z)
+    return out, out
+
+
+class Decoder(nn.Module):
+  """Transformer decoder model with optional cross-attention."""
+  emb_dim: int
+  mlp_dim: int
+  num_heads: int
+  num_layers: int
+  out_dim: int
+  seq_len: int
+  style: Literal["ar", "masked"]
+  dropout_rate: float = 0.
+  zero_embedding_init: bool = False
+
+  scan: bool = False
+  remat_policy: str = "nothing_saveable"
+
+  @nn.compact
+  def __call__(
+      self,
+      targets: jax.Array,
+      encoded: jax.Array | None = None,
+      decoder_mask: jax.Array | None = None,
+      decode: bool = False,
+      deterministic: bool = True,
+      return_reps: bool = False,
+  ) -> jax.Array | tuple[jax.Array, Mapping[str, jax.Array]]:
+    """Applies Transformer model on the inputs.
+
+    Args:
+      targets: target text tokens [B, L].
+      encoded: encoded sequence from an encoder [B, P, E].
+      decoder_mask: decoder self-attention mask.
+      decode: bool, whether to perform fast autoregressive decoding with cache.
+      deterministic: bool, deterministic or not (to apply dropout).
+      return_reps: bool, whether to return intermediate representations.
+
+    Returns:
+      output of a transformer decoder [B, L, out_dim], where out_dim is usually
+      a multiple of D.
+    """
+    if self.style == "masked" and decode:
+      raise ValueError("Cannot run masked model in cached mode!")
+
+    pos_emb = vit.get_posemb(
+        self, "learn", self.seq_len, self.emb_dim,
+        "pos_emb")
+
+    y = common.AddPositionEmbs(
+        decode=decode, name="PosEmbedTargets")(targets, pos_emb)
+
+    out = {}
+    if self.scan:
+      # Mostly followed
+      # https://github.com/google/maxtext/blob/4d99e30b3e0e0cb1d1aa11c7db7fffe18e301498/MaxText/layers.py#L1126
+      # for the scanned version.
+
+      # 1. remat
+      enc_dec_block_remat = nn.remat(
+          EncoderDecoderBlock,
+          prevent_cse=False,
+          static_argnums=(-1, -2),
+          policy=getattr(jax.checkpoint_policies, self.remat_policy, None))
+      # 2. scan
+      initializing = self.is_mutable_collection("params")
+      param_scan_axis = 1
+      params_spec = (param_scan_axis if initializing
+                     else partitioning.ScanIn(param_scan_axis))
+      dec_scanned = nn.scan(enc_dec_block_remat,
+                            variable_axes={
+                                "params": params_spec,
+                                "cache": 0,
+                            },
+                            split_rngs={"params": True, "dropout": True},
+                            in_axes=nn.broadcast,
+                            length=self.num_layers)
+      # 3. fprop
+      y, out = dec_scanned(num_heads=self.num_heads, mlp_dim=self.mlp_dim,
+                           dropout_rate=self.dropout_rate, decode=decode,
+                           name="EncDecBlock")(
+                               y, encoded, decoder_mask, deterministic)
+      # Extracting the intermediate representation from the stacked activation
+      # tensor `out`, which is a [num_layers, B, L, E] tensor. Indexing along
+      # the first axis to extract individual layers, and then averaging across
+      # the second axis, which corresponds to the sequence dimension after
+      # indexing.
+      assert out.shape[0] == self.num_layers and (
+          decode or out.shape[2] == self.seq_len), (
+              (out.shape, self.num_layers, self.seq_len))
+      out = {f"block{l}_rep": jnp.mean(out[l], axis=1)
+             for l in range(self.num_layers)}
+    else:
+      for lyr in range(self.num_layers):
+        y, _ = EncoderDecoderBlock(
+            num_heads=self.num_heads, mlp_dim=self.mlp_dim,
+            dropout_rate=self.dropout_rate, decode=decode,
+            name=f"EncDecBlock{lyr}")(y, encoded, decoder_mask=decoder_mask,
+                                      deterministic=deterministic)
+        out[f"block{lyr}_rep"] = jnp.mean(y, axis=1)
+    y = nn.LayerNorm(name="LayerNorm")(y)
+    out["pre_logits"] = jnp.mean(y, axis=1)
+
+    logits = nn.Dense(
+        self.out_dim,
+        kernel_init=nn.initializers.zeros,
+        name="LogitsDense",
+    )(y)
+    out["logits"] = logits
+    if return_reps:
+      return logits, out
+    return logits
+
+
+class Model(nn.Module):
+  """GIVT model supporting decoder-only and encoder-decoder applications."""
+  num_heads: int = 8
+  # num_layers = 0 means no encoder
+  num_layers: int = 0
+  num_decoder_layers: int = 6
+  mlp_dim: int = 2048
+  enc_dropout_rate: float = 0.
+  dec_dropout_rate: float = 0.
+  # Decoder params:
+  emb_dim: int = 512
+  num_labels: Optional[int] = 1000
+  seq_len: int = 256
+  # Encoder params:
+  patches: Sequence[int] = (16, 16)
+  input_size: Sequence[int] = (256, 256)
+  posemb_type: Literal["learn", "sincos2d"] = "learn"
+  zero_decoder_seq: bool = False
+  style: Literal["ar", "masked"] = "ar"
+
+  zero_embedding_init: bool = False
+
+  num_mixtures: int = 4
+  multivariate: bool = False
+  out_dim: int = 32
+  scale_tol: float = 1e-6
+
+  # Mask specific params.
+  mask_schedule_train: str = "cosine"
+  # Results in at least 40% masked tokens with cosine.
+  min_masking_rate_training: float = 0.3
+
+  # How to fuse mask at input:
+  # - replace: replace token[masked] with lookup(MASK)
+  # - concat: replace token[mask] with lookup(REPLACE) and concat either
+  #   lookup(NOMASK) or lookup(MASK).
+  mask_style: str = "replace"
+
+  # Set to >0 for CFG support.
+  drop_labels_probability: float = 0.0
+
+  fix_square_plus: bool = False
+
+  # If True, and mixture >1, create a GMM per channel. Otherwise, create
+  # a GMM of `dim`-dimensional Gaussians.
+  per_channel_mixtures: bool = True
+
+  scan: bool = False
+  remat_policy: str = "nothing_saveable"
+
+  @property
+  def has_encoder(self) -> bool:
+    return self.num_layers > 0
+
+  @property
+  def num_logits(self) -> int:
+    if self.multivariate:
+      assert self.num_mixtures == 1
+      # d**2 covariance, d means.
+      # Note: `round` makes pytype happy.
+      return round(self.out_dim ** 2) + self.out_dim
+
+    elif self.per_channel_mixtures:
+      # One (mu, sigma, pi) per output dimension and mixture component.
+      # Note that we predict a distribution for each output dimensions in
+      # parallel.
+      return 3 * self.num_mixtures * self.out_dim
+
+    else:
+      # Mixture weights plus mean/scale per mixture
+      return self.num_mixtures + 2 * self.num_mixtures * self.out_dim
+
+  def setup(self) -> None:
+    assert self.posemb_type == "learn"
+    assert self.num_mixtures > 0
+
+    if self.multivariate and self.num_mixtures != 1:
+      raise ValueError("Cannot do multivariate GMM!")
+
+    if self.num_layers > 0:
+      grid_size = np.array(self.input_size) // np.array(self.patches)
+
+      self.pos_emb_for_encoder = vit.get_posemb(
+          self, self.posemb_type, grid_size, self.emb_dim,
+          "pos_embedding_encoder")
+
+      self.conv = nn.Conv(self.emb_dim, self.patches, padding="VALID",
+                          strides=self.patches, name="EmbedPatches")
+
+      self.encoder = vit.Encoder(
+          depth=self.num_layers,
+          mlp_dim=self.mlp_dim,
+          num_heads=self.num_heads,
+          dropout=self.enc_dropout_rate,
+          scan=self.scan,
+          remat_policy=self.remat_policy,)
+    else:
+      self.encoder = None
+
+    # Iterator that will lead free label IDs.
+    next_label = itertools.count(self.num_labels or 0)
+    special_labels = {}
+
+    if self.style == "ar":
+      pass
+    elif self.style == "masked":
+      if self.mask_style == "replace":
+        special_labels = {_SpecialLabel.MASK: next(next_label)}
+      elif self.mask_style == "concat":
+        special_labels = {
+            _SpecialLabel.MASK: next(next_label),
+            _SpecialLabel.NOMASK: next(next_label),
+            _SpecialLabel.REPLACE: next(next_label),
+        }
+      else:
+        raise NotImplementedError(self.mask_style)
+    else:
+      raise NotImplementedError(self.style)
+
+    if self.drop_labels_probability > 0:
+      special_labels[_SpecialLabel.NOLABEL] = next(next_label)
+
+    self.special_labels = special_labels
+    lookup_size = (self.num_labels or 1) + len(self.special_labels)
+
+    self.labels_emb = nn.Embed(
+        lookup_size,
+        self.emb_dim,
+        name="EmbedLabels",
+        embedding_init=nn.initializers.zeros
+        if self.zero_embedding_init
+        else nn.initializers.normal(stddev=1.0),
+    )
+
+    self.targets_emb = nn.Dense(self.emb_dim, name="EmbedTargets")
+
+    self.decoder = Decoder(
+        num_layers=self.num_decoder_layers or self.num_layers,
+        mlp_dim=self.mlp_dim,
+        num_heads=self.num_heads,
+        out_dim=self.num_logits,
+        # In masked mode, we run with 1 more token at the input.
+        seq_len=self.seq_len + int(self.style == "masked"),
+        dropout_rate=self.dec_dropout_rate,
+        emb_dim=self.emb_dim,
+        zero_embedding_init=self.zero_embedding_init,
+        style=self.style,
+        scan=self.scan,
+        remat_policy=self.remat_policy,
+    )
+
+  def encode(self, image: jax.Array, train: bool = False) -> jax.Array:
+    """Encodes input image or embeddings."""
+    emb = self.conv(image)
+    patch_embeddings = einops.rearrange(emb, "B PH PW E -> B (PH PW) E")
+    encoded, _ = self.encoder(
+        patch_embeddings + self.pos_emb_for_encoder, deterministic=not train)
+    return encoded
+
+  def embed_labels(
+      self,
+      labels: jax.Array | None = None,
+      batch_size: int | None = None,
+  ) -> jax.Array:
+    if labels is not None:
+      # Embed class label, add a sequence dim (output shape (B, 1, E))
+      return self.labels_emb(labels)[:, None, :]
+
+    assert ((self.num_labels == 1 or self.num_labels is None)
+            and batch_size is not None)
+    # Create [BOS] token embedding
+    return self.labels_emb(jnp.zeros((batch_size,), jnp.int32))[:, None, :]
+
+  def prefill(
+      self, labels=None, batch_size=None, encoded=None, drop_labels=None
+  ):
+    labels = self._drop_labels(drop_labels, labels)
+    labels_for_prefill = self.embed_labels(labels=labels, batch_size=batch_size)
+    return self.decoder(
+        labels_for_prefill,
+        encoded=encoded,
+        decode=True)
+
+  def _decode_ar(
+      self,
+      targets: jax.Array,
+      labels: jax.Array | None = None,
+      encoded: jax.Array | None = None,
+      decode: bool = False,
+      train: bool = False,
+  ) -> tuple[jax.Array, Mapping[str, jax.Array]]:
+    """Autoregressive decoding."""
+    targets_embedded = self.targets_emb(targets)
+
+    if decode:
+      decoder_mask = None
+    else:
+      decoder_mask = nn.make_causal_mask(targets[:, :, 0])
+      b = targets.shape[0]
+      labels_embedded = self.embed_labels(labels, b)
+      assert labels_embedded.shape == (b, 1, self.emb_dim), (
+          labels_embedded.shape, (b, 1, self.emb_dim))
+      targets_embedded = jnp.concatenate(
+          [labels_embedded, targets_embedded[:, : -1]], axis=1)
+
+    logits, out = self.decoder(
+        targets_embedded,
+        encoded=encoded,
+        decoder_mask=decoder_mask,
+        decode=decode,
+        deterministic=not train,
+        return_reps=True)
+
+    return logits, out
+
+  def _get_special_label(self, size, label: _SpecialLabel):
+    return self.labels_emb(
+        jnp.full(size, self.special_labels[label], jnp.int32)
+    )
+
+  def _decode_masked(
+      self,
+      targets,
+      input_mask,
+      labels=None,
+      encoded=None,
+      train=False,
+  ):
+    """Masked decoding."""
+    b, s, _ = targets.shape
+    assert input_mask.shape == (b, s)
+
+    if self.mask_style == "replace":
+      targets_embedded = jnp.where(
+          input_mask[:, :, None],
+          self._get_special_label((b, s), _SpecialLabel.MASK),
+          self.targets_emb(targets),
+      )
+    elif self.mask_style == "concat":
+      masks = jnp.where(
+          input_mask[:, :, None],
+          self._get_special_label((b, s), _SpecialLabel.MASK),
+          self._get_special_label((b, s), _SpecialLabel.NOMASK),
+      )
+      embedded_targets = self.targets_emb(targets)
+      targets_embedded = jnp.where(
+          input_mask[:, :, None],
+          self._get_special_label((b, s), _SpecialLabel.REPLACE),
+          embedded_targets,
+      )
+      # Only take half of each to get the right embedding size.
+      targets_embedded = jnp.concatenate(
+          [masks[..., ::2], targets_embedded[..., ::2]], axis=-1
+      )
+    else:
+      raise ValueError(self.mask_style)
+
+    labels_embedded = self.embed_labels(labels, b)
+    assert labels_embedded.shape == (b, 1, self.emb_dim)
+    # Note that we do not truncate the input here, so this has shape
+    # (B, L+1, E).
+    targets_embedded = jnp.concatenate(
+        [labels_embedded, targets_embedded], axis=1)
+
+    logits = self.decoder(
+        targets_embedded,
+        encoded=encoded,
+        decoder_mask=None,
+        decode=False,
+        deterministic=not train)
+
+    logits = logits[:, 1:, ...]  # Remove class label
+    assert logits.shape[:2] == (b, s)
+    return logits
+
+  def _drop_labels(self, drop_labels_mask, labels):
+    if labels is None:
+      return None
+    if self.drop_labels_probability >= 0.999:
+      logging.warning("Dropping all labels...")
+      return jnp.full_like(labels, self.special_labels[_SpecialLabel.NOLABEL])
+    if drop_labels_mask is None:
+      return labels
+    assert _SpecialLabel.NOLABEL in self.special_labels
+    nolabel = jnp.full_like(
+        labels, self.special_labels[_SpecialLabel.NOLABEL]
+    )
+    return jnp.where(drop_labels_mask, nolabel, labels)
+
+  def decode(
+      self,
+      targets: jax.Array,
+      labels: jax.Array | None = None,
+      encoded: jax.Array | None = None,
+      decode: bool = False,
+      train: bool = False,
+      max_decode_length: int | None = None,
+      input_mask: jax.Array | None = None,
+      drop_labels: jax.Array | None = None,
+      return_reps: bool = False,
+  ) -> jax.Array | tuple[jax.Array, Mapping[str, jax.Array]]:
+    """Applies Transformer decoder-branch on encoded-input and target.
+
+    Args:
+      targets: target text tokens [B, L, out_dim].
+      labels: optional class labes, [B].
+      encoded: encoded image patches from encoder [B, P, E].
+      decode: whether to prepare and use an autoregressive cache.
+      train: whether it is training.
+      max_decode_length: optional max length for positional embeddings.
+      input_mask: If given, mask input. Required for style=="masked".
+        Shape [B, L], bool tensor. True means the token will be removed
+        from the input.
+      drop_labels: Drop labels at corresponding locations [B].
+      return_reps: whether to return intermediate representations.
+
+    Returns:
+      logits array from transformer decoder [B, L, 3 * num_mixtures * out_dim].
+    """
+    del max_decode_length
+    labels = self._drop_labels(drop_labels, labels)
+    if self.style == "ar":
+      logits, out = self._decode_ar(
+          targets, labels, encoded, decode, train)
+      if return_reps:
+        return logits, out
+      return logits
+    elif self.style == "masked":
+      assert not decode  # Cache not supported.
+      assert input_mask is not None
+      assert not return_reps  # Not implemented.
+      return self._decode_masked(targets, input_mask, labels, encoded, train)
+    else:
+      raise NotImplementedError(self.style)
+
+  def _square_plus(self, x):
+    # Via https://twitter.com/jon_barron/status/1387167648669048833
+    if self.fix_square_plus:
+      return (x + jnp.sqrt(jnp.square(x) + 4)) / 2
+    else:
+      return x + jnp.sqrt(jnp.square(x) + 4) / 2
+
+  def get_pdf(
+      self,
+      logits: jax.Array,
+      temperature_scales: float | None = None,
+      temperature_probs: float | None = None,
+  ) -> distrax.Distribution:
+    assert logits.shape[-1] == self.num_logits
+    if self.multivariate:
+      scales = logits[..., :self.out_dim ** 2]
+      locs = logits[..., self.out_dim ** 2:]
+      assert locs.shape[-1] == self.out_dim
+      scales = self._square_plus(scales)
+      # Turn into a square matrix.
+      *leading, _ = scales.shape
+      scales = scales.reshape(*leading, self.out_dim, self.out_dim)
+      # Make sure the diagonals are non zero.
+      diag_scale_tol = jnp.eye(self.out_dim) * self.scale_tol
+      scales = jnp.maximum(scales, diag_scale_tol)
+      if (t := temperature_scales) is not None:
+        scales = scales * t
+
+      # Note that there is `tfd.MultivariateNormalFullCovariance`` but it just
+      # calls linalg.cholesky on the covariance and then uses the
+      # MultivariateNormalTri class.  Using ... direcly avoids having to
+      # construct a hermetian matrix.
+      #
+      # Note that only the lower triag part of `scales` is used by applying
+      # jnp.tril. The other elements are replaced with zeros.
+      #
+      # Note on output shapes:
+      # - .sample() -> shape (..., seq_len, out_dim)
+      # - .prob()   -> shape (..., seq_len).
+      return distrax.MultivariateNormalTri(locs, scales)
+
+    elif self.per_channel_mixtures:
+      # [..., 3 * num_mixtures * out_dim] -> [..., 3 * out_dim, num_mixtures]
+      logits = jnp.reshape(logits, logits.shape[: -1] + (-1, self.num_mixtures))
+      # 3 tensors with shape [..., out_dim, num_mixtures]
+      probs, locs, scales = jnp.split(logits, 3, axis=-2)
+      if (t := temperature_probs) is not None:
+        probs = probs * t
+
+      # normalize mixture probabilities
+      probs = nn.softmax(probs)
+      scales = self._square_plus(scales)
+      # threshold scale
+      scales = jnp.maximum(scales, self.scale_tol)
+      if (t := temperature_scales) is not None:
+        scales = scales * t
+
+      # Note on output shapes:
+      # - .sample() -> shape (..., seq_len, out_dim)
+      # - .prob()   -> shape (..., seq_len, out_dim).
+      return distrax.MixtureSameFamily(
+          mixture_distribution=distrax.Categorical(probs=probs),
+          components_distribution=distrax.Normal(loc=locs, scale=scales),
+      )
+    else:
+      *shape, num_logits = logits.shape
+      assert num_logits == self.num_logits, (num_logits, self.num_logits)
+      prob_logits, other_logits = (
+          logits[..., : self.num_mixtures],
+          logits[..., self.num_mixtures :],
+      )
+      if (t := temperature_probs) is not None:
+        prob_logits = prob_logits * t
+      other_logits = jnp.reshape(
+          other_logits, (*shape, self.num_mixtures, 2, self.out_dim)
+      )
+      locs = other_logits[..., 0, :]
+      scales = self._square_plus(other_logits[..., 1, :])
+
+      scales = jnp.maximum(scales, self.scale_tol)  # Threshold scale
+      if (t := temperature_scales) is not None:
+        scales = scales * t
+
+      # prob_logits has shape (b, seq_len, m)
+      # locs/scales has shape (b, seq_len, m, d)
+      assert prob_logits.ndim == locs.ndim - 1, (prob_logits.shape, locs.shape)
+      assert locs.shape == scales.shape, (locs.shape, scales.shape)
+
+      # Note on output shapes:
+      # - .sample() -> shape (..., seq_len, out_dim)
+      # - .prob()   -> shape (..., seq_len,)
+      # - .nll()   -> shape (..., seq_len,)
+      return distrax.MixtureSameFamily(
+          mixture_distribution=distrax.Categorical(logits=prob_logits),
+          components_distribution=distrax.MultivariateNormalDiag(
+              loc=locs, scale_diag=scales
+          ),
+      )
+
+  def __call__(
+      self,
+      sequence: jax.Array,
+      labels: jax.Array | None = None,
+      *,
+      image: jax.Array | None = None,
+      decode: bool = False,
+      input_mask: jax.Array | None = None,
+      drop_labels: jax.Array | None = None,
+      train: bool = False,
+  ) -> tuple[jax.Array, distrax.Distribution]:
+    """Applies Transformer model on the inputs.
+
+    Args:
+      sequence: batch of sequences [B, L].
+      labels: class labels for class conditional generation [B].
+      image: batch of images [B, H, W, 3].
+      decode: whether to prepare and use an autoregressive cache.
+      input_mask: If given, mask input. Required for style=="masked" [B, L].
+      drop_labels: If given, drop labels of the corresponding batches [B].
+      train: whether it is training.
+
+    Returns:
+      logits array from full transformer [B, L, out_dim].
+    """
+    if self.style == "masked" and input_mask is None:
+      raise ValueError("Cannot run masked model without input mask!")
+
+    if self.encoder is not None:
+      assert image is not None
+      encoded = self.encode(image, train=train)
+    else:
+      assert image is None
+      encoded = None
+
+    logits = self.decode(sequence, labels=labels, encoded=encoded,
+                         decode=decode, input_mask=input_mask, train=train)
+    pdf = self.get_pdf(logits)
+    return logits, pdf
+
+  def get_input_mask_training(
+      self,
+      rng: jax.Array,
+      shape: tuple[int, int],
+  ) -> jax.Array | None:
+    """Creates a random maask of shape (B, L) for training masked models."""
+    if self.style == "ar":
+      return None
+    b, s = shape
+    # Sample b values in [0, 1-min_mask_ratio].
+    keep = jax.random.uniform(
+        rng, shape=(b,), maxval=1.0 - self.min_masking_rate_training
+    )
+    mask_ratio = apply_mask_schedule(keep, self.mask_schedule_train)
+    return _random_mask_with_ratios(rng, ratios=mask_ratio, seq_len=s)
+
+  def get_input_mask_teacher_forced(
+      self,
+      shape: tuple[int, int],
+  ) -> jax.Array | None:
+    """Creates a random maask of shape (B, L) for training masked models."""
+    if self.style == "ar":
+      return None
+    return jnp.zeros(shape, dtype=jnp.bool_)
+
+  def get_drop_labels(
+      self,
+      rng: jax.Array,
+      batch_size: int,
+  ) -> jax.Array | None:
+    if (p := self.drop_labels_probability) > 0:
+      return jax.random.uniform(rng, shape=(batch_size,)) <= p
+    else:
+      return None
+
+
+def load(
+    init_params: Any,
+    init_files: str | Mapping[str, str],
+    model_params: Any = None,
+    dont_load: Sequence[str] = (),
+    resample_encoder_posemb: bool = False,
+    trim_decoder_posemb: bool = False,
+) -> Any:
+  """Loads params from init checkpoint and merges into init_params."""
+  del model_params
+  if isinstance(init_files, str):
+    ckpt_params = utils.load_params(init_files)
+    ckpt_params = common.merge_params(ckpt_params, init_params, dont_load)
+
+    if resample_encoder_posemb:
+      if init_params and "pos_embedding_encoder" in init_params:
+        ckpt_params["pos_embedding_encoder"] = vit.resample_posemb(
+            old=ckpt_params["pos_embedding_encoder"],
+            new=init_params["pos_embedding_encoder"])
+
+    if trim_decoder_posemb:
+      if init_params and "pos_embedding_decoder" in init_params:
+        ckpt_params["pos_embedding_decoder"] = (
+            ckpt_params["pos_embedding_decoder"][
+                :, :init_params["pos_embedding_decoder"].shape[1], :])
+
+  else:
+    init_files = {**init_files}  # Shallow copy because we'll pop stuff off.
+
+    enc_init = init_files.pop("encoder", None)
+    if enc_init:
+      ckpt_params = init_params.copy()
+      vit_params = {
+          "pos_embedding": ckpt_params["pos_embedding_encoder"],
+          "Transformer": ckpt_params["encoder"],
+          "embedding": ckpt_params["EmbedPatches"],
+      }
+      encoder_params = vit.load(
+          vit_params, enc_init, model_cfg={},
+          dont_load=dont_load)
+      ckpt_params["encoder"] = encoder_params["Transformer"]
+      ckpt_params["pos_embedding_encoder"] = encoder_params["pos_embedding"]
+      ckpt_params["EmbedPatches"] = encoder_params["embedding"]
+    else:
+      raise ValueError("Only encoder init is supported: {}.".format(init_files))
+
+  return ckpt_params

Tipsomaly/model/big_vision/models/proj/givt/givt_test.py

@@ -0,0 +1,124 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for GIVT model."""
+
+from absl.testing import parameterized
+from big_vision.models.proj.givt import givt
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from absl.testing import absltest
+
+
+_BATCH_SIZE = 2
+_OUT_DIM = 4
+_SEQ_LEN = 16
+_NUM_MIXTURES = 4
+
+
+def _make_test_model(**overwrites):
+  config = dict(
+      num_heads=2,
+      num_decoder_layers=1,
+      mlp_dim=64,
+      emb_dim=16,
+      seq_len=_SEQ_LEN,
+      out_dim=_OUT_DIM,
+      num_mixtures=_NUM_MIXTURES,
+  )
+  config.update(overwrites)
+  return givt.Model(**config)
+
+
+class MaskedTransformerTest(parameterized.TestCase):
+
+  @parameterized.product(rng_seed=[0])
+  def test_masks(self, rng_seed):
+    m = _make_test_model(style="masked")
+    mask = m.get_input_mask_training(jax.random.PRNGKey(rng_seed), (2, 16))
+    self.assertEqual(mask.shape, (2, 16))
+    # At least one should definitly be masked out.
+    self.assertTrue(np.all(mask.sum(-1) > 1))
+
+  @parameterized.product(
+      train=[True, False],
+      multivariate=[True, False],
+      per_channel_mixtures=[True, False],
+      drop_labels_probability=[0.0, 0.1],
+      style=["masked", "ar"],
+  )
+  def test_apply(
+      self,
+      train,
+      multivariate,
+      per_channel_mixtures,
+      drop_labels_probability,
+      style,
+  ):
+    if per_channel_mixtures and multivariate:
+      self.skipTest("Not supported")
+    model = _make_test_model(
+        style=style,
+        multivariate=multivariate,
+        num_mixtures=1 if multivariate else _NUM_MIXTURES,
+        per_channel_mixtures=per_channel_mixtures,
+        drop_labels_probability=drop_labels_probability,
+    )
+    sequence = jax.random.uniform(
+        jax.random.PRNGKey(0), (_BATCH_SIZE, _SEQ_LEN, _OUT_DIM)
+    )
+    labels = jax.random.uniform(
+        jax.random.PRNGKey(0), (_BATCH_SIZE,), maxval=10
+    ).astype(jnp.int32)
+    input_mask = jax.random.uniform(
+        jax.random.PRNGKey(0), (_BATCH_SIZE, _SEQ_LEN)
+    ).astype(jnp.bool_)
+    variables = model.init(
+        jax.random.PRNGKey(0),
+        sequence,
+        labels,
+        input_mask=input_mask,
+        train=train,
+    )
+    logits, pdf = model.apply(
+        variables, sequence, labels, input_mask=input_mask, train=train
+    )
+    nll = -pdf.log_prob(sequence)
+    self.assertFalse(np.any(np.isnan(nll)))
+    if multivariate:
+      self.assertEqual(
+          logits.shape, (_BATCH_SIZE, _SEQ_LEN, _OUT_DIM**2 + _OUT_DIM)
+      )
+      self.assertEqual(nll.shape, (_BATCH_SIZE, _SEQ_LEN))
+    elif per_channel_mixtures:
+      self.assertEqual(
+          logits.shape,
+          (_BATCH_SIZE, _SEQ_LEN, 3 * _NUM_MIXTURES * _OUT_DIM),
+      )
+      self.assertEqual(nll.shape, (_BATCH_SIZE, _SEQ_LEN, _OUT_DIM))
+    else:
+      self.assertEqual(
+          logits.shape,
+          (_BATCH_SIZE, _SEQ_LEN, _NUM_MIXTURES + _NUM_MIXTURES * _OUT_DIM * 2),
+      )
+      self.assertEqual(nll.shape, (_BATCH_SIZE, _SEQ_LEN))
+
+    sample = pdf.sample(seed=jax.random.PRNGKey(0))
+    self.assertEqual(sample.shape, (_BATCH_SIZE, _SEQ_LEN, _OUT_DIM))
+
+
+if __name__ == "__main__":
+  googletest.main()

Tipsomaly/model/big_vision/models/proj/givt/parallel_decode.py

@@ -0,0 +1,523 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Decode autoregressive/bidirectional masked transformers.
+
+
+Currently, we implement MaskGIT style temperature sampling:
+
+In each step:
+1. Get P = model(inputs), predicted GMMs
+2. Get samples = sample_from(P)
+3. Get probs = P[samples], ie, model evaluated at samples.
+   We use this now as a confidence metric, but we scale the probs:
+4. probs = probs ^ 1/choice_temperature
+4. set probs[already_uncovered_points] = inf, ie, we will always keep
+   uncovered points (no resampling!)
+5. Now pick top K points from probs to keep for the next steps, where
+   K = some monotonically increasing ratio of points as we go along decoding
+"""
+
+import dataclasses
+from typing import Literal
+
+from absl import logging
+from big_vision.models.proj.givt import givt
+import distrax
+import flax
+import jax
+import jax.numpy as jnp
+
+
+_CONFIDENCE_OF_KNOWN_TOKENS = jnp.inf
+
+
+@jax.vmap
+def _get_per_batch_mask(arr, k):
+  (d,) = arr.shape
+  indices = jnp.argsort(arr)
+  valid_indices = jnp.arange(d) < k
+  return jnp.zeros((d,), jnp.bool_).at[indices].set(valid_indices)
+
+
+def _get_bottom_k_mask(arr, k):
+  *leading, d = arr.shape
+  arr = arr.reshape((-1, d))
+  mask = _get_per_batch_mask(arr, k)
+  return mask.reshape(*leading, -1)
+
+
+def mask_by_random_topk(rng, mask_len, probs, temperature=1.0):
+  """Create a mask.
+
+  Adaption of jax.random.choice where probabilities are changed by scaling with
+  `temperature` (probs = probs ^ (1/temperature)).
+
+  Additionally, this function returns a mask of tokens to mask out, which
+  are picked to be the low confidence ones. Thus, this function is roughly
+  equivalent to (but not exactly at edge cases such as prob = inf..):
+
+    keep = jax.random.choice(
+        rng, seq_len,
+        shape=(seq_len - mask_len,),
+        # NOTE: probabilities are updated with `temperature`.
+        p=jnp.power(probs, 1/temperature),
+        replace=False
+    )
+    mask = jnp.ones((seq_len,), dtype=jnp.bool_)
+    return mask.at[..., keep].set(False)
+
+  Args:
+    rng: a PRNG key used as the random key.
+    mask_len: the number to mask.
+    probs: the probabilities associated with each entry.
+    temperature: when temperature = 1.0, it's identical to jax's implementation.
+      The larger this value is, the more random the masking is picked.
+
+  Returns:
+    A binary masking map [batch_size, seq_len]. Contains True where we should
+    mask (at mask_len locations), and False where we should keep.
+  """
+  confidence = jnp.log(probs) + temperature * jax.random.gumbel(
+      rng, probs.shape)
+  return _get_bottom_k_mask(confidence, mask_len)
+
+
+@flax.struct.dataclass
+class DecodeState:
+  """Holds decoding state data."""
+
+  rng: jax.Array  # Sampling random state.
+  # The position of the decoding loop in the length dimension. Scalar int32.
+  step: jax.Array
+  # What we input at each step. Starts from all masks and is uncovered by
+  # sampling. Note that this is an array with leading
+  # dimension `num_steps + 1` because we start with all masked tokens and then
+  # need `num_steps` to uncover all, i.e., the final output is given by
+  # all_inputs_q[-1, ...].
+  all_inputs_q: jax.Array  # float32 [num_steps + 1, batch, seq_len, c]
+  # Has a 1 for every _uncovered_ point.
+  uncovered_per_step: jax.Array  # bool_ [num_steps, batch, seq_len]
+  logits_per_step: jax.Array  # [num_steps, batch, seq_len, num_logits]
+  uncond_logits_per_step: jax.Array  # [num_steps, batch, seq_len, num_logits]
+  prob_per_step: jax.Array  # Probability per step.
+  # If CFG: Rejection sampling success rate.
+  rejection_sampling_success_per_step: jax.Array
+
+  @classmethod
+  def make(
+      cls,
+      initial_rng: jax.Array,
+      all_masked_input: jax.Array,
+      num_logits: int,
+      num_steps: int,
+  ) -> "DecodeState":
+    """Creates the initial state."""
+    b, seq_len, c = all_masked_input.shape
+    all_inputs_q = jnp.broadcast_to(
+        all_masked_input,
+        (num_steps + 1, b, seq_len, c),
+    )
+    return cls(
+        initial_rng,
+        step=jnp.array(0),
+        all_inputs_q=all_inputs_q,
+        uncovered_per_step=jnp.full((num_steps, b, seq_len), False, jnp.bool_),
+        logits_per_step=jnp.full(
+            (num_steps, b, seq_len, num_logits), jnp.nan, jnp.float32
+        ),
+        uncond_logits_per_step=jnp.full(
+            (num_steps, b, seq_len, num_logits), jnp.nan, jnp.float32
+        ),
+        prob_per_step=jnp.full((num_steps, b, seq_len), jnp.nan, jnp.float32),
+        rejection_sampling_success_per_step=jnp.full(
+            (num_steps,), jnp.nan, jnp.float32
+        ),
+    )
+
+  @property
+  def current_inputs_q(self) -> jax.Array:
+    """Returns the current quantized input."""
+    return self.all_inputs_q[self.step, ...]
+
+  @property
+  def num_steps(self) -> int:
+    """Returns number of decode steps."""
+    return self.uncovered_per_step.shape[0]
+
+  def _steps_mask(self) -> jax.Array:
+    return jnp.arange(self.num_steps) <= self.step
+
+  @property
+  def total_uncovered(self) -> jax.Array:
+    """Returns the total uncovered mask up to and including current step."""
+    return self.uncovered_per_step.sum(
+        axis=0, where=self._steps_mask()[:, jnp.newaxis, jnp.newaxis]
+    ).astype(jnp.bool_)
+
+  def split_rng(self) -> tuple["DecodeState", jax.Array]:
+    """Splits of RNG for the current step."""
+    rng, step_rng = jax.random.split(self.rng, 2)
+    return self.replace(rng=rng), step_rng
+
+  def set_next_input(self, next_input_q: jax.Array) -> "DecodeState":
+    """Sets the input for the next step."""
+    return self._set_row("all_inputs_q", self.step + 1, next_input_q)
+
+  def set_uncover_at_current_step(self, uncovered: jax.Array) -> "DecodeState":
+    """Sets what was uncovered after the current step."""
+    return self._set_row("uncovered_per_step", self.step, uncovered)
+
+  def set_logits_at_current_step(self, logits: jax.Array) -> "DecodeState":
+    return self._set_row("logits_per_step", self.step, logits)
+
+  def set_uncond_logits_at_current_step(
+      self, logits: jax.Array
+  ) -> "DecodeState":
+    return self._set_row("uncond_logits_per_step", self.step, logits)
+
+  def set_rejection_sampling_success_at_current_step(
+      self, success: jax.Array
+  ) -> "DecodeState":
+    return self._set_row(
+        "rejection_sampling_success_per_step", self.step, success
+    )
+
+  def set_prob_at_current_step(self, prob: jax.Array) -> "DecodeState":
+    return self._set_row("prob_per_step", self.step, prob)
+
+  def increment_step(self) -> "DecodeState":
+    """Increments step."""
+    return self.replace(step=self.step + 1)
+
+  def _set_row(self, attr_name, row_index, row_value):
+    """Sets one row of the variables that have shape (num_steps, ...)."""
+    current_value = getattr(self, attr_name)
+    _, *expected_shape = current_value.shape
+    if row_value.shape != tuple(expected_shape):
+      raise ValueError(f"Expected {row_value.shape} == {expected_shape}!")
+    if row_value.dtype != current_value.dtype:
+      raise ValueError(f"Expected {row_value.dtype} == {current_value.dtype}")
+    new_value = current_value.at[row_index, ...].set(row_value)
+    return self.replace(**{attr_name: new_value})
+
+
+@dataclasses.dataclass(frozen=True)
+class MaskedGenerationConfig:
+  """Config for masked generation.
+
+  Attributes:
+    num_steps: Number of sampling steps.
+    should_anneal_temperature: If given, anneal choice temperature as we go
+      through the sampling steps.
+    choice_temperature: Temperature for picking points.
+    ordering: How to order to select. Supports:
+      maskgit: Maskgit style, use P[samples]
+    schedule: Inference mask schedule.
+    cfg_inference_weight: CFG Inference weight.
+  """
+  num_steps: int = 16
+  should_anneal_temperature: bool = True
+  choice_temperature: float = 1.0
+  ordering: Literal["maskgit"] = "maskgit"
+  schedule: str = "cosine"
+  cfg_inference_weight: float = 0.0
+
+
+def _assert_single_component_get_loc_scale(
+    pdf: distrax.Distribution, rng=None, mixture=None
+):
+  """Extracts loc and scale from a single mixture GMM."""
+  if not isinstance(pdf, distrax.MixtureSameFamily):
+    raise ValueError(f"Expected mixture! Got {type(pdf)}")
+  components_d = pdf.components_distribution
+  if isinstance(components_d, distrax.MultivariateNormalDiag):
+    loc, scale_diag = components_d.loc, components_d.scale_diag
+    b, s, m, _ = loc.shape
+    if mixture is None:
+      assert rng is not None
+      # Shape (b, seq)
+      mixture = pdf.mixture_distribution.sample(seed=rng)
+      mixture = jax.nn.one_hot(mixture, num_classes=m, axis=-1)
+    assert mixture.shape == (b, s, m), (mixture.shape, loc.shape)
+    loc = (loc * mixture[..., None]).sum(-2)
+    scale_diag = (scale_diag * mixture[..., None]).sum(-2)
+    return loc, scale_diag, mixture
+  else:
+    loc, scale = components_d.loc, components_d.scale
+    if loc.shape[-1] != 1 or scale.shape[-1] != 1:
+      raise ValueError(f"Expected one mixture! {loc.shape}/{scale.shape}")
+    return loc[..., 0], scale[..., 0], None
+
+
+class CFGDensity:
+  """Helper to get probability and samples via CFG."""
+
+  pdf_c: distrax.Distribution
+  pdf_u: distrax.Distribution
+  w: float
+  simple: distrax.Distribution
+  fac: jax.Array
+
+  def __init__(
+      self,
+      pdf_c: distrax.Distribution,
+      pdf_u: distrax.Distribution,
+      w: float,
+      rng: jax.Array,
+  ) -> None:
+    loc_c, scale_c, mixture = _assert_single_component_get_loc_scale(pdf_c, rng)
+    # Note: RNG only needed when we have mixtures, to select components.
+    loc_u, scale_u, _ = _assert_single_component_get_loc_scale(
+        pdf_u, rng, mixture=mixture
+    )
+
+    # Definitly wider than whatever we had before. The mean should be slightly
+    # away though!
+    loc_simple = loc_c
+    scale_simple = jnp.stack([scale_c, scale_u], -1).max(-1) * 2
+    self.simple = distrax.Normal(loc_simple, scale_simple)
+
+    self.pdf_c = distrax.Normal(loc_c, scale_c)
+    self.pdf_u = distrax.Normal(loc_u, scale_u)
+    self.w = w
+
+    assert loc_c.ndim == 3, loc_c.shape
+    points = loc_c[jnp.newaxis, ...] + jnp.linspace(-10, 10, 1001).reshape(
+        -1, 1, 1, 1
+    )
+    p_at_c, _ = self._unnormalized_p(points)
+
+    self.fac = jnp.max(p_at_c / self.simple.prob(loc_c), axis=0)
+    jax.debug.print("🎲 CFG {fac}", fac=self.fac.mean())
+
+  def _unnormalized_p(self, x):
+    w = self.w
+    logp_cfg = (1 + w) * self.pdf_c.log_prob(x) - w * self.pdf_u.log_prob(x)
+    return jnp.exp(logp_cfg), logp_cfg
+
+  def rejection_sample(
+      self,
+      seed: jax.Array,
+      max_samples: int = 1_000,
+  ) -> tuple[jax.Array, jax.Array]:
+    """Rejection sampling, try `max_samples`, take first match."""
+    rng_sample, rng_uni = jax.random.split(seed, 2)
+    # Shape (max_samples, b, seq_len, c)
+    xs = self.simple.sample(seed=rng_sample, sample_shape=(max_samples,))
+    facq = self.fac * self.simple.prob(xs)
+    ys = jax.random.uniform(rng_uni, shape=facq.shape, minval=0.0, maxval=facq)
+    # Shape (max_samples, b, seq_len, c), True where `xs` is a valid sample
+    # from p. We might have anywhere between 0 and `max_samples` valid samples!
+    p, _ = self._unnormalized_p(xs)
+    mask = ys < p
+    # Now we need to do fancy tricks to get the first element in `mask` that is
+    # True. We do this by making a shifted mask that is False for every element
+    # after the first True.
+    # > Example:
+    # mask            [0, 1, 0, 1, 0, 0, 1, 0]
+    # > implies:
+    # cmask           [0, 1, 1, 1, 1, 1, 1, 1]
+    # shifted_cmask   [0, 0, 1, 1, 1, 1, 1, 1]
+    # keep            [0, 1, 0, 0, 0, 0, 0, 0]  # <- picks the first valid!
+    cmask = jnp.cumsum(mask, axis=0).astype(jnp.bool_)
+    shifted_cmask = jnp.pad(
+        cmask, [(1, 0), (0, 0), (0, 0), (0, 0)], constant_values=False
+    )[:-1]
+    assert shifted_cmask.shape == mask.shape
+    keep = jnp.logical_and(cmask, jnp.logical_not(shifted_cmask))
+    # Now we can grab the first valid sample by doing a sum over the
+    # `max_samples` dimension.
+    sample = jnp.where(keep, xs, 0).sum(0)
+    # If the rejection sampler fails, we fall back to the conditional
+    # distribution.
+    ok = mask.sum(0) > 0  # Shape (b, seq_len, c)
+    # jax.debug.print("🎲 CFG ok {ok}%", ok=ok.mean() * 100)
+    sample = jnp.where(
+        ok, sample, self.pdf_c.sample(seed=rng_sample)
+    )
+    return sample, ok.mean() * 100
+
+  def sample(
+      self,
+      seed: jax.Array,
+      max_samples: int = 1_000,
+  ) -> jax.Array:
+    result, ok = self.rejection_sample(seed, max_samples)
+    jax.debug.print("Debug ok={ok}%", ok=ok)
+    return result
+
+  # Unnormalized! But we only use it for ordering.
+  def prob(self, xs: jax.Array) -> jax.Array:
+    p, _ = self._unnormalized_p(xs)
+    return p
+
+  def log_prob(self, xs: jax.Array) -> jax.Array:
+    _, lp = self._unnormalized_p(xs)
+    return lp
+
+
+def decode_masked(
+    rng: jax.Array,
+    labels: jax.Array,
+    seq_len: int,
+    feature_dim: int,
+    model: givt.Model,
+    variables: flax.core.FrozenDict,
+    config: MaskedGenerationConfig,
+) -> DecodeState:
+  """Implements an masked bidirectional sampling loop.
+
+  This function implements the loop from the docstring.
+
+  Args:
+    rng: RNG, only required if sampling.
+    labels: Shape (b,), labels per batch. Determines batch size.
+    seq_len: How many tokens to sample per batch.
+    feature_dim: Output dimension of the VAE, i.e., number of channels, `c`.
+    model: GIVT model to sample from.
+    variables: Variables of the model.
+    config: Configures style.
+
+  Returns:
+    Final state.
+  """
+  logging.info("Masked Generation Config:\n%s", config)
+
+  if model.style != "masked":
+    raise ValueError(f"Need masked model! Got `{model.style}`.")
+
+  (b,) = labels.shape
+  all_masked_input = jnp.zeros((b, seq_len, feature_dim))
+  init_state = DecodeState.make(
+      rng,
+      all_masked_input,
+      num_logits=model.num_logits,
+      num_steps=config.num_steps,
+  )
+
+  def loop_cond_fn(state: DecodeState):
+    return state.step < state.num_steps
+
+  def tokens_to_logits(tokens, input_mask, drop_labels=None):
+    return model.apply(
+        variables,
+        tokens,
+        labels=labels,
+        # Note that the model applies the mask token internally given the input.
+        input_mask=input_mask,
+        drop_labels=drop_labels,
+        method="decode",
+    )
+
+  def loop_body_fn(state: DecodeState) -> DecodeState:
+    # 1 where we should mask, cumulative.
+    unknown = jnp.logical_not(state.total_uncovered)
+
+    # Defines the mask ratio for the next round. The number to mask out is
+    # determined by mask_ratio * unknown_number_in_the_beginning.
+    ratio = (state.step + 1) / config.num_steps
+    # Note that the mask schedule inverts the function, so `mask_ratio` givts
+    # near 1 and goes to 0 monotonically.
+    mask_ratio = givt.apply_mask_schedule(ratio, method=config.schedule)
+    mask_len = jnp.floor(seq_len * mask_ratio).reshape(1, 1)
+    num_unknown = jnp.sum(unknown, axis=-1, keepdims=True)
+    mask_len = jnp.maximum(
+        0,
+        # Keeps at least one of prediction in this round: Avoids the case where
+        # mask_len is equal to num_unknown, in which case the mask is not
+        # updated! We substract 1 to always remove at least one masked token.
+        jnp.minimum(num_unknown - 1, mask_len))
+
+    # Run model ---
+    logits = tokens_to_logits(state.current_inputs_q, unknown)
+    # Book keeping: store all logits.
+    state = state.set_logits_at_current_step(logits)
+
+    pdf = model.get_pdf(logits)
+    state, sample_rng = state.split_rng()
+    if config.cfg_inference_weight > 0:
+      drop_all_labels = jnp.full((b,), True, jnp.bool_)
+      logits_uncond = tokens_to_logits(
+          state.current_inputs_q, unknown, drop_labels=drop_all_labels
+      )
+      state = state.set_uncond_logits_at_current_step(logits_uncond)
+      pdf_uncond = model.get_pdf(logits_uncond)
+      state, cfg_rng = state.split_rng()
+      pdf = CFGDensity(
+          pdf_c=pdf,
+          pdf_u=pdf_uncond,
+          w=config.cfg_inference_weight,
+          rng=cfg_rng,
+      )
+      sample, rejection_sampling_success = pdf.rejection_sample(sample_rng)
+      state = state.set_rejection_sampling_success_at_current_step(
+          rejection_sampling_success
+      )
+    else:
+      sample = pdf.sample(seed=sample_rng)
+
+    # Sample at the unknown spots.
+    sampled = jnp.where(unknown[:, :, None], sample, state.current_inputs_q)
+    assert sampled.shape == (b, seq_len, feature_dim), (
+        sampled.shape,
+        b,
+        seq_len,
+        feature_dim,
+    )
+
+    prob = pdf.prob(sampled)
+    if model.multivariate:
+      assert prob.ndim == 2  # (b, seq_len) already
+    elif model.per_channel_mixtures or config.cfg_inference_weight > 0:
+      # Independence accross channels.
+      # This reduction is also required when using CFG and also
+      # `model.per_channel_mixtures == False` due to the 2-step CFG redefining
+      # the pdf, but the reduction is not needed without CFG.
+      prob = prob.prod(-1)
+    state = state.set_prob_at_current_step(prob)
+
+    if config.ordering == "maskgit":
+      ordering = jnp.where(unknown, prob, _CONFIDENCE_OF_KNOWN_TOKENS)
+    else:
+      raise NotImplementedError(config.ordering)
+
+    assert ordering.shape == (b, seq_len), (ordering.shape, b, seq_len)
+
+    temp = config.choice_temperature
+    if config.should_anneal_temperature:
+      temp *= (1. - ratio)
+
+    # True where we should mask input. Note that this is cumulative (ie this
+    # starts with all True and keeps getting more False entries as we go through
+    # the steps).
+    state, choice_rng = state.split_rng()
+    masking = mask_by_random_topk(choice_rng, mask_len, ordering, temp)
+    assert masking.shape == (b, seq_len)
+    masking = jnp.where(mask_len == 0, jnp.zeros_like(masking), masking)
+
+    # Remove the masked tokens from the sampled array for safety (the model will
+    # again apply the mask anyway...).
+    sampled = jnp.where(masking[:, :, None], jnp.zeros_like(sampled), sampled)
+
+    # Get next_uncover ---
+    # New tokens to uncover (non cumulative): where it was unknown
+    # but is now known.
+    next_uncover = jnp.logical_and(unknown, jnp.logical_not(masking))
+    assert next_uncover.shape == (b, seq_len), (next_uncover.shape, b, seq_len)
+    state = state.set_uncover_at_current_step(next_uncover)
+    state = state.set_next_input(sampled)
+    return state.increment_step()
+
+  return jax.lax.while_loop(loop_cond_fn, loop_body_fn, init_state)

Tipsomaly/model/big_vision/models/proj/givt/parallel_decode_test.py

@@ -0,0 +1,154 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from absl.testing import parameterized
+from big_vision.models.proj.givt import givt
+from big_vision.models.proj.givt import parallel_decode
+import chex
+import jax
+import jax.numpy as jnp
+
+from absl.testing import absltest
+
+
+_BATCH_SIZE = 2
+_OUT_DIM = 4
+_SEQ_LEN = 6
+_NUM_MIXTURES = 4
+
+
+def _make_test_model(**overwrites):
+  config = dict(
+      num_heads=2,
+      num_decoder_layers=1,
+      mlp_dim=64,
+      emb_dim=16,
+      seq_len=_SEQ_LEN,
+      out_dim=_OUT_DIM,
+      num_mixtures=_NUM_MIXTURES,
+      style="masked",
+  )
+  config.update(overwrites)
+  return givt.Model(**config)
+
+
+def _mask(*flags):
+  return jnp.asarray(flags).astype(jnp.bool_)
+
+
+class HelperTest(googletest.TestCase):
+
+  def test_get_first_n(self):
+    with self.subTest("ordered"):
+      values = jnp.asarray([4, 3, 2, 1, 0])
+      k = jnp.asarray([3], jnp.int32)
+      chex.assert_trees_all_equal(
+          parallel_decode._get_bottom_k_mask(values, k), _mask(0, 0, 1, 1, 1)
+      )
+
+    with self.subTest("equal_values"):
+      values = jnp.ones((5,))
+      k = jnp.asarray([3], jnp.int32)
+      chex.assert_trees_all_equal(
+          parallel_decode._get_bottom_k_mask(values, k), _mask(1, 1, 1, 0, 0)
+      )
+
+    with self.subTest("equal_values"):
+      values = jnp.asarray([1, 2, 2, 2, 3])
+      k = jnp.asarray([3], jnp.int32)
+      chex.assert_trees_all_equal(
+          parallel_decode._get_bottom_k_mask(values, k), _mask(1, 1, 1, 0, 0)
+      )
+
+
+class ParallelDecodeTest(parameterized.TestCase):
+
+  def _make_model(self, **overwrites):
+    model = _make_test_model(**overwrites)
+    sequence = jax.random.uniform(
+        jax.random.PRNGKey(0), (_BATCH_SIZE, _SEQ_LEN, _OUT_DIM)
+    )
+    labels = jax.random.uniform(
+        jax.random.PRNGKey(0), (_BATCH_SIZE,), maxval=10
+    ).astype(jnp.int32)
+    input_mask = jax.random.uniform(
+        jax.random.PRNGKey(0), (_BATCH_SIZE, _SEQ_LEN)
+    ).astype(jnp.bool_)
+    variables = model.init(
+        jax.random.PRNGKey(0),
+        sequence,
+        labels,
+        input_mask=input_mask,
+        train=False,
+    )
+    return model, variables
+
+  def _test_model(self, rng, model, variables, config):
+    labels = jnp.ones((_BATCH_SIZE,), dtype=jnp.int32)
+    state = parallel_decode.decode_masked(
+        rng,
+        seq_len=_SEQ_LEN,
+        feature_dim=_OUT_DIM,
+        labels=labels,
+        model=model,
+        variables=variables,
+        config=config,
+    )
+    self.assertEqual(int(state.step), 4)
+    # Each point uncovered exactly once.
+    chex.assert_trees_all_equal(
+        state.uncovered_per_step.sum(0),
+        jnp.ones((_BATCH_SIZE, _SEQ_LEN), dtype=jnp.int32),
+    )
+
+  @parameterized.product(
+      rng_seed=[1, 2],
+      choice_temperature=[1.0, 4.0],
+      multivariate=[True, False],
+  )
+  def test_decode_masked(self, rng_seed, choice_temperature, multivariate):
+    rng = jax.random.PRNGKey(rng_seed)
+    model, variables = self._make_model(
+        num_mixtures=1 if multivariate else _NUM_MIXTURES,
+        multivariate=multivariate,
+    )
+    config = parallel_decode.MaskedGenerationConfig(
+        num_steps=4,
+        choice_temperature=choice_temperature,
+    )
+    self._test_model(rng, model, variables, config)
+
+  @parameterized.product(
+      rng_seed=[1, 2],
+      choice_temperature=[1.0, 4.0],
+      w=[0.0, 1.0, 3.0],
+      per_channel_mixtures=[True, False],
+  )
+  def test_cfg(self, rng_seed, choice_temperature, w, per_channel_mixtures):
+    rng = jax.random.PRNGKey(rng_seed)
+    model, variables = self._make_model(
+        num_mixtures=1 if per_channel_mixtures else 3,
+        drop_labels_probability=0.1,
+        per_channel_mixtures=per_channel_mixtures,
+    )
+    config = parallel_decode.MaskedGenerationConfig(
+        num_steps=4,
+        choice_temperature=choice_temperature,
+        cfg_inference_weight=w,
+    )
+    self._test_model(rng, model, variables, config)
+
+
+if __name__ == "__main__":
+  googletest.main()

Tipsomaly/model/big_vision/models/proj/givt/vae.py

@@ -0,0 +1,94 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Abstract VAE model class.
+
+Gaussian encoder and decoder (the latter assumed to have constant variance).
+
+Inspiration drawn from https://github.com/pytorch/examples/tree/main/vae.
+"""
+
+import abc
+from typing import Optional, Mapping
+
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+
+
+class Model(nn.Module, metaclass=abc.ABCMeta):
+  """Abstract VAE model class."""
+
+  codeword_dim: Optional[int] = None
+  code_len: int = 256
+  code_dropout: str = "none"
+
+  @abc.abstractmethod
+  def encode(
+      self,
+      x: jax.Array,
+      *,
+      train: bool = False,
+  ) -> tuple[jax.Array, jax.Array]:
+    ...
+
+  def reparametrize(
+      self,
+      mu: jax.Array,
+      logvar: jax.Array,
+      rng: jax.Array | None = None,
+  ) -> jax.Array:
+    std = jnp.exp(0.5 * logvar)
+    if rng is None:
+      rng = self.make_rng("dropout")
+    eps = jax.random.normal(rng, shape=std.shape, dtype=std.dtype)
+    return mu + std * eps
+
+  @abc.abstractmethod
+  def decode(
+      self, x: jax.Array,
+      train: bool = False,
+  ) -> jax.Array | Mapping[str, jax.Array]:
+    ...
+
+  def code_dropout_fn(self, z: jax.Array, *, train: bool = False) -> jax.Array:
+    # "seq" drops out tokens later in the sequence with higher probablility than
+    # tokens earlier in the sequence.
+    assert self.code_dropout in ["none", "seq", "random"]
+    if train and self.code_dropout != "none":
+      importance = jnp.linspace(1.0, 0.0, self.code_len + 2)[1:-1]
+      thr = jax.random.uniform(self.make_rng("dropout"), z.shape[:1])
+      mask = importance[None, :] > thr[:, None]
+      if self.code_dropout == "random":
+        mask = jax.random.permutation(
+            self.make_rng("dropout"), mask, axis=-1, independent=True)
+      z = z * mask[:, :, None]
+    return z
+
+  def __call__(
+      self,
+      x: jax.Array,
+      *,
+      train: bool = False,
+  ) -> tuple[jax.Array | Mapping[str, jax.Array], Mapping[str, jax.Array]]:
+    mu, logvar = self.encode(x, train=train)
+    # Only reparametrize when training for simplicity.
+    if train:
+      z = self.reparametrize(mu, logvar)
+    else:
+      z = mu
+    z = self.code_dropout_fn(z, train=train)
+    x = self.decode(z, train=train)
+    return x, {"mu": mu, "logvar": logvar, "z": z}

Tipsomaly/model/big_vision/models/proj/givt/vit.py

@@ -0,0 +1,188 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Simple VAE fork of the UViM VQ-VAE (proj/uvim/vit.py) with small changes."""
+
+from typing import Optional, Sequence, Mapping, Any
+
+from big_vision import utils
+from big_vision.models import common
+from big_vision.models import vit
+from big_vision.models.proj.givt import vae
+
+import einops
+import flax.linen as nn
+import flax.training.checkpoints
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+class Model(vae.Model):
+  """ViT model."""
+
+  input_size: Sequence[int] = (256, 256)
+  patch_size: Sequence[int] = (16, 16)
+  width: int = 768
+  enc_depth: int = 6
+  dec_depth: int = 6
+  mlp_dim: Optional[int] = None
+  num_heads: int = 12
+  posemb: str = "learn"  # Can also be "sincos2d"
+  dropout: float = 0.0
+  head_zeroinit: bool = True
+  bottleneck_resize: bool = False
+  inout_specs: Optional[Mapping[str, tuple[int, int]]] = None
+  scan: bool = False
+  remat_policy: str = "nothing_saveable"
+
+  def setup(self) -> None:
+    self.grid_size = np.array(self.input_size) // np.array(self.patch_size)
+
+    self.embedding = nn.Conv(
+        self.width, self.patch_size, strides=self.patch_size,
+        padding="VALID", name="embedding")
+
+    self.pos_embedding_encoder = vit.get_posemb(
+        self, self.posemb, self.grid_size, self.width, "pos_embedding_encoder")
+    self.encoder = vit.Encoder(
+        depth=self.enc_depth,
+        mlp_dim=self.mlp_dim,
+        num_heads=self.num_heads,
+        dropout=self.dropout,
+        scan=self.scan,
+        remat_policy=self.remat_policy,
+        name="encoder")
+
+    if not self.bottleneck_resize:
+      self.bottleneck_downsample = self.param(
+          "bottleneck_downsample",
+          nn.initializers.xavier_uniform(),
+          (np.prod(self.grid_size), self.code_len))
+
+    if not self.bottleneck_resize:
+      self.bottleneck_upsample = self.param(
+          "bottleneck_upsample",
+          nn.initializers.xavier_uniform(),
+          (self.code_len, np.prod(self.grid_size)))
+
+    self.pos_embedding_decoder = vit.get_posemb(
+        self, self.posemb, self.grid_size, self.width, "pos_embedding_decoder")
+    self.decoder = vit.Encoder(
+        depth=self.dec_depth,
+        mlp_dim=self.mlp_dim,
+        num_heads=self.num_heads,
+        dropout=self.dropout,
+        scan=self.scan,
+        remat_policy=self.remat_policy,
+        name="decoder")
+
+    # Setting num_outputs to 2 * codeword_dim to predict mean and variance per
+    # element
+    self.encoder_head = nn.Dense(self.codeword_dim * 2 or self.width * 2)
+    self.decoder_stem = nn.Dense(self.width)
+
+    kw = {"kernel_init": nn.initializers.zeros} if self.head_zeroinit else {}
+
+    if self.inout_specs is not None:
+      num_out_channels = sum(
+          num_classes for _, num_classes in self.inout_specs.values())
+    else:
+      num_out_channels = 3
+
+    self.head = nn.Dense(
+        num_out_channels * np.prod(self.patch_size),
+        name="decoder_head", **kw)
+
+  def encode(
+      self,
+      x: jax.Array,
+      *,
+      train: bool = False,
+  ) -> tuple[jax.Array, jax.Array]:
+    if self.inout_specs is not None:
+      one_hot_inputs = []
+      for in_ch, num_classes in self.inout_specs.values():
+        one_hot_inputs.append(nn.one_hot(x[..., in_ch], num_classes))
+      x = jnp.concatenate(one_hot_inputs, axis=-1)
+    x = self.embedding(x)
+    x = einops.rearrange(x, "b h w c -> b (h w) c")
+
+    x, _ = self.encoder(x + self.pos_embedding_encoder, deterministic=not train)
+
+    if self.bottleneck_resize:
+      x = einops.rearrange(x, "b (h w) c -> b h w c",
+                           h=self.grid_size[0], w=self.grid_size[1])
+      l = int(np.round(self.code_len ** 0.5))
+      x = jax.image.resize(
+          x, (x.shape[0], l, l, x.shape[3]),
+          method="linear")
+      x = einops.rearrange(x, "b h w c -> b (h w) c")
+    else:
+      x = jnp.einsum("btc,tn->bnc", x, self.bottleneck_downsample)
+
+    x = self.encoder_head(x)
+
+    mu, logvar = jnp.split(x, 2, axis=-1)
+    return mu, logvar
+
+  def decode(
+      self,
+      x: jax.Array,
+      train: bool = False,
+  ) -> jax.Array | Mapping[str, jax.Array]:
+    x = self.decoder_stem(x)
+
+    if self.bottleneck_resize:
+      l = int(np.round(self.code_len ** 0.5))
+      x = einops.rearrange(x, "b (h w) c -> b h w c", h=l, w=l)
+      x = jax.image.resize(
+          x, (x.shape[0], self.grid_size[0], self.grid_size[1], x.shape[3]),
+          method="linear")
+      x = einops.rearrange(x, "b h w c -> b (h w) c")
+    else:
+      x = jnp.einsum("bnc,nt->btc", x, self.bottleneck_upsample)
+
+    x, _ = self.decoder(x + self.pos_embedding_decoder, deterministic=not train)
+    x = self.head(x)
+    # c = 3 for RGB images
+    x = einops.rearrange(x, "b (h w) (p q c) -> b (h p) (w q) c",
+                         h=self.grid_size[0], w=self.grid_size[1],
+                         p=self.patch_size[0], q=self.patch_size[1])
+
+    if self.inout_specs is None:
+      x = jnp.clip(x, -1.0, 1.0)
+    else:
+      x_dict = {}
+      channel_index = 0
+      for name, (_, num_channels) in self.inout_specs.items():
+        x_dict[name] = x[..., channel_index : channel_index + num_channels]
+        channel_index += num_channels
+      x = x_dict
+
+    return x
+
+
+def load(
+    init_params: Any,
+    init_file: str,
+    model_params: Any = None,
+    dont_load: Sequence[str] = (),
+) -> Any:
+  """Loads params from init checkpoint and merges into init_params."""
+  del model_params
+  params = flax.core.unfreeze(utils.load_params(init_file))
+  if init_params is not None:
+    params = common.merge_params(params, init_params, dont_load)
+  return params

Tipsomaly/model/big_vision/models/proj/image_text/naflex_vit.py

@@ -0,0 +1,300 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""NaFlex ViT = NaViT + FlexiViT.
+
+Based on:
+* FlexiViT: https://arxiv.org/abs/2212.08013
+* NaViT: https://arxiv.org/abs/2307.06304
+"""
+
+import re
+from big_vision.models import vit
+import big_vision.models.proj.image_text.utils as it_utils
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+def _decode_posemb(posemb):
+  if (m := re.fullmatch(r"learn_2d(\(\d+\))", posemb)):
+    grid_size = int(m.groups()[0][1:-1])
+    return "learn_2d", grid_size
+  return posemb, None
+
+
+def _pos_emb_resize(pos_emb, shapes, coords, l):
+  """Resizes the positional embeddings to match the input image size.
+  
+  Args:
+    pos_emb: Positional embeddings.
+    shapes: Image shapes (usually `coords.max(axis=1) + 1`).
+    coords: Patch coordinates.
+    l: Maximum number of patches per side. Necesary in order to have a static
+      return shape.
+
+  Setting l to 64 is a heuristic. Ideally, we would use
+  `l = tokens.shape[1]` here, but that requires too much memory,
+  especially for high-resolution inputs. Using a lower value
+  effectively limits the maximum resolution to `l x patch_size`.
+  Resolutions above that will lead to NaNs in the positional
+  embeddings and NaN model outputs.
+  Note: this value can be adjusted post-hoc without retraining.
+
+  Returns:
+    Postional embeddings for every patch.
+  """
+
+  def resize_fn(shape, coords):
+    emb = jax.image.scale_and_translate(
+        pos_emb,
+        shape=(l, l, pos_emb.shape[-1]),
+        spatial_dims=(0, 1),
+        scale=shape / jnp.asarray(pos_emb.shape[:2]),
+        translation=jnp.asarray([0, 0]),
+        method="bilinear", antialias=True)
+    gather_dim = jax.lax.GatherDimensionNumbers(
+        offset_dims=(1,),
+        collapsed_slice_dims=(0, 1),
+        start_index_map=(0, 1, 2)
+    )
+    return jax.lax.gather(
+        emb,
+        jnp.pad(coords, [[0, 0], [0, 1]]),
+        gather_dim,
+        [1, 1, emb.shape[-1]],
+        mode="fill")
+  return it_utils.batch_shmap(
+      jax.vmap(resize_fn, in_axes=(0, 0), out_axes=0),
+      shapes, coords)
+
+
+class Encoder1DBlock(nn.Module):
+  """Single transformer encoder block (MHSA + MLP)."""
+  mlp_dim: int | None = None  # Defaults to 4x input dim
+  num_heads: int = 12
+  dropout: float = 0.0
+  dtype_mm: str = "float32"
+
+  @nn.compact
+  def __call__(self, x, mask=None, deterministic=True):
+    if mask is not None:
+      mask = mask[..., None, :, :]  # Broadcast mask along the head dim.
+
+    out = {}
+    x = nn.with_logical_constraint(x, ("act_batch", "act_len", "act_emb"))
+    y = nn.LayerNorm()(x)
+    y = out["sa"] = nn.MultiHeadDotProductAttention(
+        num_heads=self.num_heads,
+        kernel_init=nn.initializers.xavier_uniform(),
+        deterministic=deterministic,
+        dtype=self.dtype_mm,
+    )(y, y, mask=mask)
+    y = nn.with_logical_constraint(y, ("act_batch", "act_len", "act_emb"))
+    y = nn.Dropout(rate=self.dropout)(y, deterministic)
+    x = out["+sa"] = x + y
+
+    y = nn.LayerNorm()(x)
+    y = out["mlp"] = vit.MlpBlock(
+        mlp_dim=self.mlp_dim, dropout=self.dropout,
+        dtype_mm=self.dtype_mm,
+    )(y, deterministic)
+    y = nn.with_logical_constraint(y, ("act_batch", "act_len", "act_emb"))
+    y = nn.Dropout(rate=self.dropout)(y, deterministic)
+    x = out["+mlp"] = x + y
+    x = nn.with_logical_constraint(x, ("act_batch", "act_len", "act_emb"))
+    return x, out
+
+
+class Encoder(nn.Module):
+  """Transformer Model Encoder for sequence to sequence translation."""
+  depth: int
+  mlp_dim: int | None = None  # Defaults to 4x input dim
+  num_heads: int = 12
+  dropout: float = 0.0
+  scan: bool = False
+  remat_policy: str = "nothing_saveable"
+  dtype_mm: str = "float32"
+
+  @nn.compact
+  def __call__(self, x, mask=None, deterministic=True):
+    out = {}
+
+    if self.scan:
+      block = nn.remat(
+          Encoder1DBlock,
+          prevent_cse=False,
+          static_argnums=(3,),  # 0=self, 3=deterministic
+          policy=getattr(jax.checkpoint_policies, self.remat_policy, None),
+          )
+      x, scan_out = nn.scan(
+          block,
+          variable_axes={"params": 0},
+          split_rngs={"params": True, "dropout": True},
+          in_axes=nn.broadcast,
+          length=self.depth)(
+              name="encoderblock",
+              dtype_mm=self.dtype_mm,
+              mlp_dim=self.mlp_dim,
+              num_heads=self.num_heads,
+              dropout=self.dropout)(x, mask, deterministic)
+      for lyr in range(self.depth):
+        out[f"block{lyr:02d}"] = jax.tree.map(lambda o, l=lyr: o[l], scan_out)
+    else:
+      # Input Encoder
+      for lyr in range(self.depth):
+        block_cur = Encoder1DBlock(
+            name=f"encoderblock_{lyr}",
+            dtype_mm=self.dtype_mm,
+            mlp_dim=self.mlp_dim, num_heads=self.num_heads,
+            dropout=self.dropout)
+        x, out[f"block{lyr:02d}"] = block_cur(x, mask, deterministic)
+      out["pre_ln"] = x  # Alias for last block, but without the number in it.
+
+    return nn.LayerNorm(name="encoder_norm")(x), out
+
+
+class MAPHead(nn.Module):
+  """Multihead Attention Pooling."""
+  mlp_dim: int | None = None  # Defaults to 4x input dim
+  num_heads: int = 12
+
+  @nn.compact
+  def __call__(self, x, mask=None):
+    n, l, d = x.shape  # pylint: disable=unused-variable
+    probe = self.param("probe", nn.initializers.xavier_uniform(),
+                       (1, 1, d), x.dtype)
+    probe = jnp.tile(probe, [n, 1, 1])
+
+    if mask is not None:
+      mask = mask[..., None, None, :]  # Add query and head dims.
+
+    x = nn.MultiHeadDotProductAttention(
+        num_heads=self.num_heads,
+        kernel_init=nn.initializers.xavier_uniform())(probe, x, mask=mask)
+
+    y = nn.LayerNorm()(x)
+    x = x + vit.MlpBlock(mlp_dim=self.mlp_dim)(y)
+    return x[:, 0]
+
+
+class _Model(nn.Module):
+  """ViT model."""
+
+  num_classes: int | None = None
+  width: int = 768
+  depth: int = 12
+  mlp_dim: int | None = None  # Defaults to 4x input dim
+  num_heads: int = 12
+  rep_size: int | bool = False
+  pool_type: str = "gap"  # Can also be "map" or "tok"
+  head_zeroinit: bool = True
+  scan: bool = False
+  # or "dots_with_no_batch_dims_saveable" for more speed (memory costly)
+  remat_policy: str = "nothing_saveable"
+  dtype_mm: str = "float32"
+
+  posemb: str = "learn_2d(64)"
+  nposemb: int | None = None  # Needs to be overwritten
+
+  patchln_pre: bool = False
+  patchln_post: bool = False
+
+  @nn.compact
+  def __call__(self, image, *, train=False):
+    out = {}
+
+    patches, ptype, yabs, xabs = image
+    patches = jnp.asarray(patches, self.dtype_mm)  # BN(hw3) of float32
+
+    if self.patchln_pre:
+      patches = nn.LayerNorm(name="patchln_pre")(patches)
+
+    # Embed the patches.
+    tokens = out["stem"] = nn.Dense(
+        self.width, name="embedding", dtype=self.dtype_mm)(patches)
+
+    if self.patchln_post:
+      tokens = nn.LayerNorm(name="patchln_post")(tokens)
+
+    x = tokens
+    posemb, posemb_grid_size = _decode_posemb(self.posemb)
+    if posemb == "learn_2d":
+      posembs = self.param(
+          "pos_embedding",
+          nn.initializers.normal(stddev=1/np.sqrt(self.width)),
+          (self.nposemb, self.nposemb, self.width), self.dtype_mm)
+      coords = jnp.stack([yabs, xabs], axis=-1)
+      shapes = coords.max(axis=1) + 1
+      # See comment in `_pos_emb_resize` for details.
+      x += _pos_emb_resize(posembs, shapes, coords, posemb_grid_size or 64)
+    else:
+      raise ValueError(f"Unknown posemb: '{self.posemb}'")
+
+    out["with_posemb"] = x
+
+    # Only use patch tokens in self-attention:
+    sa_mask = ptype == 1  # 1 == patch (pad is 0).
+    sa_mask = jnp.logical_and(sa_mask[..., :, None], sa_mask[..., None, :])
+    x, out["encoder"] = Encoder(
+        depth=self.depth,
+        mlp_dim=self.mlp_dim,
+        num_heads=self.num_heads,
+        scan=self.scan,
+        remat_policy=self.remat_policy,
+        dtype_mm=self.dtype_mm,
+        name="Transformer")(
+            x, mask=sa_mask, deterministic=not train)
+    out["encoded"] = x
+
+    # Ignore the padding tokens when pooling:
+    pool_mask = (ptype == 1)  # 1 == patch (not pad)
+    if self.pool_type == "map":
+      maphead = MAPHead(num_heads=self.num_heads, mlp_dim=self.mlp_dim)
+      x = maphead(x, mask=pool_mask)
+    elif self.pool_type == "gap":
+      pool_mask = pool_mask[..., None]
+      x = jnp.sum(x * pool_mask, axis=1) / jnp.sum(pool_mask, axis=1)
+    elif self.pool_type == "max":
+      # Tested in (internal link)
+      pool_mask = pool_mask[..., None]
+      ignore = jnp.where(pool_mask, 0, jnp.finfo(x.dtype).min)
+      x = jnp.max(pool_mask * x + ignore, axis=1)
+    elif self.pool_type == "none":
+      pass
+    else:
+      raise ValueError(f"Unknown pool type: '{self.pool_type}'")
+    out["head_input"] = x
+
+    if self.rep_size:
+      rep_size = self.width if self.rep_size is True else self.rep_size
+      hid = nn.Dense(rep_size, name="pre_logits")
+      x = nn.tanh(hid(x))
+
+    out["pre_logits"] = x
+
+    if self.num_classes:
+      kw = {"kernel_init": nn.initializers.zeros} if self.head_zeroinit else {}
+      head = nn.Dense(self.num_classes, name="head", **kw)
+      x = out["logits"] = head(x)
+
+    return x, out
+
+
+def Model(num_classes=None, *, variant=None, **kw):  # pylint: disable=invalid-name
+  """Factory function, because linen really don't like what I'm doing!"""
+  return _Model(num_classes, **{**vit.decode_variant(variant), **kw})
+
+load = vit.load

Tipsomaly/model/big_vision/models/proj/image_text/two_towers.py

@@ -0,0 +1,178 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Transformer encoders both for text and for images."""
+
+import importlib
+from typing import Any, Optional, Tuple, Union
+from absl import logging
+
+from big_vision import utils
+import flax.linen as nn
+import jax.numpy as jnp
+
+ConfigDict = Any
+
+
+class Model(nn.Module):
+  """Two towers transformer."""
+  image: Optional[ConfigDict] = None
+  text: Optional[ConfigDict] = None
+  text_model: str = "proj.image_text.text_transformer"
+  image_model: str = "vit"
+  out_dim: Union[int, Tuple[int, int]] = 128
+  temperature_init: float = 1.0
+  bias_init: Optional[float] = None
+
+  @nn.compact
+  def __call__(self, image, text=None, learnable_prompts=None, learning_method=None, **kw):
+    """Returns (B,C) image and (B,C) text representations."""
+
+    # Support calling without text or without image, for example for few-shot.
+    ztxt, zimg = None, None
+    out = {}
+    out_dims = self.out_dim
+    if isinstance(out_dims, int):
+      out_dims = (out_dims, out_dims)
+
+    # Embed the text:
+    if text is not None:
+      text_model = importlib.import_module(
+          f"big_vision.models.{self.text_model}"
+      ).Model(**{"num_classes": out_dims[1], **(self.text or {})}, name="txt")
+
+      ztxt, out_txt = text_model(text, learnable_prompts=learnable_prompts, learning_method=learning_method, **kw)
+      for k, v in out_txt.items():
+        out[f"txt/{k}"] = v
+
+      # Normalize the embeddings the models give us.
+      out["txt/norm"] = jnp.linalg.norm(ztxt, axis=1, keepdims=True)
+      out["txt/normalized"] = ztxt = ztxt / (out["txt/norm"] + 1e-8)
+
+    if image is not None:
+      image_model = importlib.import_module(
+          f"big_vision.models.{self.image_model}"
+      ).Model(**{"num_classes": out_dims[0], **(self.image or {})}, name="img")  # pylint: disable=not-a-mapping
+
+      zimg, out_img = image_model(image, **kw)
+      for k, v in out_img.items():
+        out[f"img/{k}"] = v
+
+      # Normalize the embeddings the models give us.
+      out["img/norm"] = jnp.linalg.norm(zimg, axis=1, keepdims=True)
+      out["img/normalized"] = zimg = zimg / (out["img/norm"] + 1e-8)
+
+      out["img/2d_norm"] = jnp.linalg.norm(out_img['encoded'], axis=2, keepdims=True)
+      out["img/2d_normalized"] = out_img['encoded'] / (out["img/2d_norm"] + 1e-8)
+
+    temp_init = jnp.log(self.temperature_init)
+    t = self.param("t",
+                   lambda key, shape, dtype: temp_init * jnp.ones(shape, dtype),
+                   (1,), jnp.float32)
+    out["t"] = jnp.exp(t)
+
+    out["t/parameter"] = t
+    if (b_init := self.bias_init) is not None:
+      out["b"] = self.param("b", lambda k, s, d: b_init * jnp.ones(s, d),
+                            (1,), jnp.float32)
+
+    # We could actually play with pre-multiplying by temperature here, such
+    # that out["t"] is nothing special to the trainer anymore.
+
+    return zimg, ztxt, out
+
+def load(init_params, init_files, model_cfg, img_load_kw={}, txt_load_kw={}):  # pylint: disable=dangerous-default-value
+  """Loads both towers, `init_files` is now a dict with `img` and `txt` keys."""
+  if isinstance(init_files, str):
+    init_files = VANITY_NAMES.get(init_files, init_files)
+
+  if isinstance(init_files, str):
+    # A shortcut for a single file checkpoint of a two_towers model.
+    if "bias_init" in model_cfg.keys():
+      logging.info("loading img, txt, t, and b from a single checkpoint.")
+      init_files = {k: f"{init_files}:{k}" for k in ("img", "txt", "t", "b")}
+    else:
+      logging.info("loading img, txt, and t from a single checkpoint.")
+      init_files = {k: f"{init_files}:{k}" for k in ("img", "txt", "t")}
+  else:
+    init_files = {**init_files}  # Shallow copy because we'll pop stuff off.
+
+  if not init_params:  # Convenience to skip checks in colab.
+    init_params = {"img": None, "txt": None}
+  restored_params = {**init_params}
+
+  img_init = init_files.pop("image", init_files.pop("img", None))
+  if img_init:
+    restored_params["img"] = importlib.import_module(
+        f"big_vision.models.{model_cfg.get('image_model', 'vit')}"
+    ).load(init_params["img"], img_init, model_cfg.image, **img_load_kw)
+
+  txt_init = init_files.pop("text", init_files.pop("txt", None))
+  if txt_init:
+    restored_params["txt"] = importlib.import_module(
+        f"big_vision.models.{model_cfg.get('text_model', 'proj.image_text.text_transformer')}"  # pylint: disable=line-too-long
+    ).load(init_params["txt"], txt_init, model_cfg.text, **txt_load_kw)
+
+  t_init = init_files.pop("temperature", init_files.pop("t", None))
+  if t_init:
+    restored_params["t"] = utils.load_params(t_init)
+
+  b_init = init_files.pop("bias", init_files.pop("b", None))
+  if b_init:
+    restored_params["b"] = utils.load_params(b_init)
+
+  assert not init_files, (
+      f"There's something unused left in `config.model_init`. You probably got "
+      f"a typo. Here it is: {init_files}")
+
+  return restored_params
+
+
+# Shortcut names for some canonical paper checkpoints:
+VANITY_NAMES = {
+    # pylint: disable=line-too-long
+    # SigLIP image encoder checkpoints from https://arxiv.org/abs/2303.15343
+    "SigLIP B/16 224": "gs://big_vision/siglip/webli_en_b16_224_63724782.npz",
+    "SigLIP B/16 256": "gs://big_vision/siglip/webli_en_b16_256_60500360.npz",
+    "SigLIP B/16 384": "gs://big_vision/siglip/webli_en_b16_384_68578854.npz",
+    "SigLIP B/16 512": "gs://big_vision/siglip/webli_en_b16_512_68580893.npz",
+    "SigLIP L/16 256": "gs://big_vision/siglip/webli_en_l16_256_60552751.npz",
+    "SigLIP L/16 384": "gs://big_vision/siglip/webli_en_l16_384_63634585.npz",
+    "SigLIP So400m/14 224": "gs://big_vision/siglip/webli_en_so400m_224_57633886.npz",
+    "SigLIP So400m/14 384": "gs://big_vision/siglip/webli_en_so400m_384_58765454.npz",
+    "SigLIP B/16-i18n 256": "gs://big_vision/siglip/webli_i18n_b16_256_66117334.npz",
+
+    # SigLIP 2 image and text encoder checkpoints from https://arxiv.org/abs/2502.14786
+    "SigLIP2 B/16 224": "gs://big_vision/siglip2/siglip2_b16_224.npz",
+    "SigLIP2 B/16 256": "gs://big_vision/siglip2/siglip2_b16_256.npz",
+    "SigLIP2 B/16 384": "gs://big_vision/siglip2/siglip2_b16_384.npz",
+    "SigLIP2 B/16 512": "gs://big_vision/siglip2/siglip2_b16_512.npz",
+    "SigLIP2 B/32 256": "gs://big_vision/siglip2/siglip2_b32_256.npz",
+    "SigLIP2 L/16 256": "gs://big_vision/siglip2/siglip2_l16_256.npz",
+    "SigLIP2 L/16 384": "gs://big_vision/siglip2/siglip2_l16_384.npz",
+    "SigLIP2 L/16 512": "gs://big_vision/siglip2/siglip2_l16_512.npz",
+    "SigLIP2 So400m/14 224": "gs://big_vision/siglip2/siglip2_so400m14_224.npz",
+    "SigLIP2 So400m/14 384": "gs://big_vision/siglip2/siglip2_so400m14_384.npz",
+    "SigLIP2 So400m/16 256": "gs://big_vision/siglip2/siglip2_so400m16_256.npz",
+    "SigLIP2 So400m/16 384": "gs://big_vision/siglip2/siglip2_so400m16_384.npz",
+    "SigLIP2 So400m/16 512": "gs://big_vision/siglip2/siglip2_so400m16_512.npz",
+    "SigLIP2 g-opt/16 256": "gs://big_vision/siglip2/siglip2_g-opt16_256.npz",
+    "SigLIP2 g-opt/16 384": "gs://big_vision/siglip2/siglip2_g-opt16_384.npz",
+    # SigLIP 2 NaFlex image and text encoder checkpoints.
+    # These need `image_model="proj.image_text.naflex_vit"` for the image encoder
+    # and a non-standard preprocessing, see configs/proj/image_text/README_siglip2.md.
+    "SigLIP2 B/16 NaFlex": "gs://big_vision/siglip2/siglip2_b16_naflex.npz",
+    "SigLIP2 So400m/16 NaFlex": "gs://big_vision/siglip2/siglip2_so400m16_naflex.npz",
+    # pylint: enable=line-too-long
+}

Tipsomaly/model/big_vision/pp/proj/clippo/pp_ops.py

@@ -0,0 +1,153 @@
+# Copyright 2022 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Preprocessing functions for CLIP with Pixels Only (CLIPPO)."""
+from absl import logging
+from  big_vision.pp import utils
+from big_vision.pp.registry import Registry
+import numpy as np
+import tensorflow as tf
+
+
+@Registry.register("preprocess_ops.render_unifont")
+@utils.InKeyOutKey(indefault="texts", outdefault="image")
+def get_pp_render_text(image_size: int, font_size: int = 16, max_chars=768,
+                       background_brightness=127, text_brightness=0,
+                       lower=True, monospace=False, spacing=1, min_width=4,
+                       resize_method="area"):
+  """Renders text as image, using binary Unifont.
+
+  Largely based on Jeffrey Sorensen's text rendering implementation.
+
+  Args:
+    image_size: Width/height of output image.
+    font_size: Font size to use. Recommended to leave at 16, as this requires
+    no resizing, and is safe.
+    max_chars: Maximum inpute characters to render, to make faster.
+    background_brightness: (r, g, b) of background pixels.
+    text_brightness: (r, g, b) of text pixels.
+    lower: whether to lowercase.
+    monospace: if False, text characters are horizontally trimmed according to
+      `spacing` and `minwidth` args.
+    spacing: # pixels between each letter.
+    min_width: Minimum width of each letter. Useful to make sure e.g. spaces and
+      full stops aren't collapsed to nothing.
+    resize_method: resize method to use if fontsize != 16.
+
+  Returns:
+    Function which renders text as an image.
+  """
+  bit_embedding = np.zeros((0x200000, 32), dtype=np.uint8)
+  colpattern = {64: range(32),
+                32: sorted(tuple(range(0, 32, 4)) + tuple(range(2, 32, 4)))}
+
+  unifont_path = "big_vision/pp/proj/clippo/unifont-9.0.06.hex"
+  unifont_upper_path = "big_vision/pp/proj/clippo/unifont_upper-9.0.06.hex"
+
+  with tf.io.gfile.GFile(unifont_path) as f:
+    for line in f:
+      row = int(line[0:4], 16)
+      hexbits = line[5:-1]
+      bit_embedding[row, colpattern[len(hexbits)]] = bytearray.fromhex(hexbits)
+
+  with tf.io.gfile.GFile(unifont_upper_path) as f:
+    for line in f:
+      row = int(line[0:6], 16)
+      hexbits = line[7:-1]
+      bit_embedding[row, colpattern[len(hexbits)]] = bytearray.fromhex(hexbits)
+
+  params = tf.constant(bit_embedding, dtype=tf.uint8)
+
+  def trim_letter(letter):
+    """Remove white space based on the letter size."""
+    v = tf.reduce_max(letter, axis=0)
+    has_pixels = tf.reshape(tf.where(v), (-1,), name="RS5")
+    no_pixels = tf.equal(tf.reduce_max(v), 0)
+    first = tf.cond(no_pixels, lambda: tf.constant(0, tf.int64),
+                    lambda: has_pixels[0])
+    last = tf.cond(no_pixels, lambda: tf.constant(0, tf.int64),
+                   lambda: has_pixels[-1])
+
+    first = tf.maximum(first - spacing, 0)
+    last = tf.maximum(last + spacing, first + min_width)
+    return tf.RaggedTensor.from_tensor(tf.transpose(letter[:, first:last]))
+
+  def to_image(rendered, width, height=None):
+    """Makes a nice square image from a long string of rendered charcaters."""
+    height = height or width
+    max_letter_width = tf.reduce_max(rendered.row_lengths(1))
+    row_lengths = tf.cast(tf.cumsum(rendered.row_lengths(1)), tf.float32)
+    div = tf.cast(width - max_letter_width, tf.float32)  # For rounding errors.
+    row_idx = tf.cast(tf.floor(row_lengths / div), tf.int64)
+    row_idx = tf.RaggedTensor.from_value_rowids(tf.range(tf.shape(rendered)[0]),
+                                                row_idx)
+    trimmed = tf.gather(rendered, row_idx, axis=0)
+    trimmed = trimmed.merge_dims(1, 2)
+    trimmed = trimmed.to_tensor(default_value=0)
+    trimmed = tf.transpose(trimmed, (0, 2, 1))
+    trimmed = tf.reshape(trimmed, (-1, tf.shape(trimmed)[-1]), name="RS4")
+    trimmed = trimmed[:height]
+
+    wpad = width - tf.shape(trimmed)[1]
+    hpad = height - tf.shape(trimmed)[0]
+    padded = tf.pad(trimmed, [[0, hpad], [0, wpad]])
+    tf.assert_equal(tf.shape(padded), tf.constant((height, width)))
+    return tf.ensure_shape(padded, (width, height))
+
+  def render(text):
+    if lower:
+      text = tf.strings.lower(text)
+    text = tf.reshape(text, (-1,))[0]
+    ids = tf.strings.unicode_decode(text, "UTF-8")
+    if max_chars:
+      ids = ids[:max_chars]
+    embed = tf.nn.embedding_lookup(params, ids)  # Get the letters
+    # Each letter is 32 uint8s, but we want binary 16x16 grid.
+    # The following does that in a rather hard to parse way.
+    vertical = tf.reshape(embed, [1, -1])
+    repl = tf.reshape(tf.transpose(tf.tile(vertical, multiples=[8, 1])), [-1])
+    ones = tf.ones_like(repl)
+    index = tf.cumsum(ones, exclusive=True)
+    sevens = tf.cast(tf.fill(tf.shape(repl), 7), tf.uint8)
+    moded = tf.bitwise.bitwise_and(index, sevens)
+    shifted = tf.bitwise.right_shift(repl,
+                                     tf.bitwise.bitwise_xor(moded, sevens))
+    anded = tf.bitwise.bitwise_and(shifted, ones)
+    # And finally, letters; binary, 0 = background, 1 = letter.
+    letters = tf.reshape(anded, [tf.shape(ids)[0], 16, 16])
+
+    if font_size != 16:
+      logging.warning("The unifont text rendering function is highly optimized "
+                      "for font size 16; using font size %i might lead to "
+                      "suboptimal rendering and might degrade performance.",
+                      font_size)
+      letters = tf.image.resize(letters[..., None], (font_size, font_size),
+                                method=resize_method, antialias=True)
+      letters = tf.squeeze(letters, axis=-1)
+
+    if monospace:
+      letters = tf.RaggedTensor.from_tensor(tf.transpose(letters, (0, 2, 1)))
+    else:
+      letters = tf.RaggedTensor.from_tensor(letters)
+      signature = tf.RaggedTensorSpec(shape=(None, font_size), ragged_rank=1,
+                                      dtype=letters.dtype)
+      letters = tf.map_fn(trim_letter, letters, fn_output_signature=signature)
+
+    img = to_image(letters, image_size)[..., None]    # A nice square image.
+    img *= (text_brightness - background_brightness)  # Rescale value range.
+    img += background_brightness
+
+    return tf.image.grayscale_to_rgb(tf.cast(img, tf.uint8))
+
+  return render

Tipsomaly/model/big_vision/pp/proj/flaxformer/bert_ops.py

@@ -0,0 +1,86 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""BERT-related preprocessing ops (using WordPiece tokenizer)."""
+
+from big_vision.pp import utils
+from big_vision.pp.registry import Registry
+import tensorflow as tf
+import tensorflow_text
+
+
+# Internally using
+# BasicTokenizer
+# https://github.com/tensorflow/text/blob/df5250d6cf1069990df4bf55154867391ab5381a/tensorflow_text/python/ops/bert_tokenizer.py#L67
+# WordpieceTokenizer
+# https://github.com/tensorflow/text/blob/master/tensorflow_text/python/ops/wordpiece_tokenizer.py
+def _create_bert_tokenizer(vocab_path):
+  """Returns cls_token id and tokenizer to use in a tf.Dataset.map function."""
+  # Create tokenizer inside a tf.init_scope so the vocab is only loaded from
+  # disk once per dataset iterator (see: http://(internal link)).
+  # TODO: Make a local copy of vocab if creating many iterators.
+  with tf.init_scope():
+    tokenizer = tensorflow_text.BertTokenizer(
+        vocab_path,
+        token_out_type=tf.int32,
+        lower_case=True,
+    )
+
+  with tf.io.gfile.GFile(vocab_path) as f:
+    vocab = f.read().split("\n")
+  cls_token = vocab.index("[CLS]")
+
+  return cls_token, tokenizer
+
+
+@Registry.register("preprocess_ops.bert_tokenize")
+@utils.InKeyOutKey(indefault=None, outdefault="labels")
+def get_pp_bert_tokenize(vocab_path, max_len, sample_if_multi=True):
+  """Extracts tokens with tensorflow_text.BertTokenizer.
+
+  Args:
+    vocab_path: Path to a file containing the vocabulry for the WordPiece
+      tokenizer. It's the "vocab.txt" file in the zip file downloaded from
+      the original repo https://github.com/google-research/bert
+    max_len: Number of tokens after tokenization.
+    sample_if_multi: Whether the first text should be taken (if set to `False`),
+      or whether a random text should be tokenized.
+
+  Returns:
+    A preprocessing Op.
+  """
+
+  cls_token, tokenizer = _create_bert_tokenizer(vocab_path)
+
+  def _pp_bert_tokenize(labels):
+
+    labels = tf.reshape(labels, (-1,))
+    labels = tf.concat([labels, [""]], axis=0)
+    if sample_if_multi:
+      num_texts = tf.maximum(tf.shape(labels)[0] - 1, 1)  # Don't sample "".
+      txt = labels[tf.random.uniform([], 0, num_texts, dtype=tf.int32)]
+    else:
+      txt = labels[0]  # Always works, since we append "" earlier on.
+
+    token_ids = tokenizer.tokenize(txt[None])
+    padded_token_ids, mask = tensorflow_text.pad_model_inputs(
+        token_ids, max_len - 1)
+    del mask  # Recovered from zero padding in model.
+    count = tf.shape(padded_token_ids)[0]
+    padded_token_ids = tf.concat(
+        [tf.fill([count, 1], cls_token), padded_token_ids], axis=1)
+    return padded_token_ids[0]
+
+  return _pp_bert_tokenize
+  

Tipsomaly/model/big_vision/pp/proj/flaxformer/bert_ops_test.py

@@ -0,0 +1,69 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for bert_ops."""
+
+import tempfile
+
+from big_vision import input_pipeline
+import big_vision.pp.builder as pp_builder
+import big_vision.pp.ops_general  # pylint: disable=unused-import
+from big_vision.pp.proj.flaxformer import bert_ops  # pylint: disable=unused-import
+import tensorflow as tf
+
+
+# BERT vocabulary for testing.
+_BERT_VOCAB = [
+    "[PAD]",
+    "[UNK]",
+    "more",
+    "than",
+    "one",
+    "[CLS]",
+    "[SEP]",
+]
+
+
+def _create_ds(pp_str, tensor_slices, num_examples):
+  return input_pipeline.make_for_inference(
+      tf.data.Dataset.from_tensor_slices(tensor_slices),
+      num_ex_per_process=[num_examples],
+      preprocess_fn=pp_builder.get_preprocess_fn(pp_str),
+      batch_size=num_examples,
+  )[0]
+
+
+class BertOpsTest(tf.test.TestCase):
+
+  def test_tokenize(self):
+    inkey = "texts"
+    vocab_path = f"{tempfile.mkdtemp()}/vocab.txt"
+    with open(vocab_path, "w") as f:
+      f.write("\n".join(_BERT_VOCAB))
+    pp_str = (
+        f"bert_tokenize(inkey='{inkey}', vocab_path='{vocab_path}', max_len=5)"
+        f"|keep('labels')"
+    )
+    tensor_slices = {
+        inkey: tf.ragged.constant([["one more"], ["more than one"], [""]])
+    }
+    ds = _create_ds(pp_str, tensor_slices, 3)
+    self.assertAllEqual(
+        next(iter(ds))["labels"],
+        [[5, 4, 2, 0, 0], [5, 2, 3, 4, 0], [5, 0, 0, 0, 0]],
+    )
+
+
+if __name__ == "__main__":
+  tf.test.main()

Tipsomaly/model/big_vision/pp/proj/image_text/ops_naflex_test.py

@@ -0,0 +1,77 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for NaFlex preprocessing ops."""
+
+import copy
+
+from absl.testing import parameterized
+from big_vision.pp.proj.image_text import ops_naflex as pp
+import numpy as np
+import tensorflow as tf
+
+
+def get_image_data(h, w):
+  img = tf.random.uniform((h, w, 3), 0, 255, tf.int32)  # Can't ask uint8!?
+  return {"image": tf.cast(img, tf.uint8)}
+
+
+class NaflexTest(tf.test.TestCase, parameterized.TestCase):
+
+  def tfrun(self, ppfn, data):
+    # Run once as standalone, as could happen eg in colab.
+    yield tf.nest.map_structure(np.array, ppfn(copy.deepcopy(data)))
+
+    # And then once again as part of tfdata pipeline.
+    # You'd be surprised how much these two differ!
+    tfdata = tf.data.Dataset.from_tensors(copy.deepcopy(data))
+    for npdata in tfdata.map(ppfn).as_numpy_iterator():
+      yield npdata
+
+  @parameterized.parameters(
+      (6, 8),
+      (7, 9),
+      (8, 10),
+  )
+  def test_patchify_valid(self, h, w):
+    """Tests the patchification op."""
+    op = pp.get_patchify((3, 4))
+    inputs = get_image_data(h, w)
+    for data in self.tfrun(op, inputs):
+      self.assertEqual(data["image"]["patches"].shape, (4, 3*4*3))
+      self.assertAllEqual(
+          data["image"]["patches"][-1],
+          np.array(inputs["image"])[3:6, 4:8, :].flatten())
+      self.assertAllEqual(data["image"]["yidx"], [0, 0, 1, 1])
+      self.assertAllEqual(data["image"]["xidx"], [0, 1, 0, 1])
+
+  @parameterized.named_parameters([
+      ("square_121_exact", (48, 48), 3, 121, (33, 33)),
+      ("square_225_inexact", (112, 109), 7, 225, (105, 105)),
+      ("square_64_exact", (176, 176), 11, 64, (88, 88)),
+      ("rect_12_exact", (256, 64), 16, 12, (96, 32)),
+      ("rect_15_exact_ps8", (256, 64), 8, 15, (56, 16)),
+      ("rect_16_inexact", (63, 241), 16, 16, (32, 128)),
+      ("rect_less_than_patch", (16, 512), 16, 16, (16, 256)),
+  ])
+  def test_pp_resize_to_sequence(
+      self, image_size, patch_size, seq_len, expected_image_size):
+    """Tests the AR-preserving `resize_to_sequence` op."""
+    op = pp.get_resize_to_sequence(patch_size, seq_len)
+    inputs = get_image_data(*image_size)
+    for outputs in self.tfrun(op, inputs):
+      self.assertAllEqual(outputs["image"].shape, expected_image_size + (3,))
+
+if __name__ == "__main__":
+  tf.test.main()

Tipsomaly/model/big_vision/pp/proj/paligemma/ops.py

@@ -0,0 +1,168 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""pp ops."""
+
+import functools
+
+from big_vision.pp import ops_text
+from big_vision.pp import utils
+from big_vision.pp.registry import Registry
+import big_vision.pp.tokenizer as bv_tok
+import numpy as np
+import tensorflow as tf
+
+
+@Registry.register('tokenizers.gemma')
+def get_tokenizer_gemma(
+    tokensets=(),
+    model='gs://big_vision/gemma_tokenizer.model',
+):
+  # See (internal link) for colab playground.
+  return ops_text.SentencepieceTokenizer(model=model, tokensets=tokensets)
+
+
+@functools.cache
+def tokenize_constant(model, text, bos='no', eos='no', length=None):
+  """Tokenize a constant string, with memoization."""
+  assert eos in ('no', 'yes', 'sticky')
+  assert bos in ('no', 'yes')
+  tokenizer = bv_tok.get_tokenizer(model)
+  tokens = tokenizer.to_int(
+      text, bos=bos == 'yes', eos=eos in ('yes', 'sticky'))
+
+  if length is None:
+    return tokens
+
+  if len(tokens) > length:
+    if eos == 'sticky':
+      return np.r_[tokens[:length-1], tokens[-1]]
+    else:
+      return tokens[:length]
+  else:
+    return np.pad(tokens, [(0, length - len(tokens))],
+                  constant_values=tokenizer.pad_token)
+
+
+@Registry.register('preprocess_ops.tolen')
+@utils.InKeyOutKey(indefault=None, outdefault=None, with_data=True)
+def get_tolen(length, *, sticky_end=False, pad_value=None, pad_key=None):
+  """Gets token to a fixed length."""
+  def _tolen(x, data):
+    if not length:
+      return x
+
+    xlen = tf.shape(x)[0]
+
+    if sticky_end:
+      trunc_fn = lambda: tf.concat([x[:length - 1], x[-1:]], axis=0)
+    else:
+      trunc_fn = lambda: x[:length]
+
+    # Potentially get the pad value from a data key (to be tokenizer agnostic).
+    pad_value_ = pad_value
+    if pad_key:
+      pad_value_ = data[pad_key]
+      # If coming from a previous tokenization op, it's probably 1D; take first.
+      if getattr(pad_value_, 'ndim', 0) == 1:
+        pad_value_ = pad_value_[0]
+    assert pad_value_ is not None, 'Need either pad_value or pad_key.'
+
+    pad_fn = lambda: tf.pad(x, [(0, length - xlen)], constant_values=pad_value_)
+    out = tf.cond(xlen >= length, trunc_fn, pad_fn)
+    out.set_shape([length])
+    return out
+  return _tolen
+
+
+@Registry.register('preprocess_ops.tok')
+def get_tokenize(model, length=None, *, bos='no', eos='no',
+                 text=None, key=None, inkey=None, outkey=None):
+  """Tokenizes and optionally truncates/pads a string."""
+
+  assert eos in ('no', 'yes', 'sticky')
+  assert bos in ('no', 'yes')
+  outkey_ = outkey or key
+  inkey_ = inkey or key
+
+  if text is not None:
+    assert inkey is None, 'Either inkey or text, not both.'
+    tokens = tokenize_constant(model, text, bos=bos, eos=eos, length=length)
+    tokens = tf.cast(tokens, tf.int32)  # Same dtype as in-graph tokenizer.
+    def _pp_tokenize_text(data):
+      data[outkey_] = tokens
+      return data
+    return _pp_tokenize_text
+
+  tokenizer = bv_tok.get_tokenizer(model)
+
+  def _pp_tokenize(data):
+    assert getattr(data[inkey_], 'ndim', 0) == 0, (
+        f'Can only tokenize single string ({inkey_}, {data[inkey_].ndim}-D)')
+
+    toks = tokenizer.to_int_tf_op(
+        data[inkey_], bos=bos == 'yes', eos=eos in ('yes', 'sticky'))
+    toks = tf.ensure_shape(toks, [None])
+    tolen = get_tolen(
+        length, sticky_end=eos == 'sticky',
+        pad_value=bv_tok.get_tokenizer(model).pad_token,
+        key='tmp',
+    )
+    toks = tolen({'tmp': toks})['tmp']
+
+    data[outkey_] = toks
+    return data
+  return _pp_tokenize
+
+
+@Registry.register('preprocess_ops.masked_concat')
+def get_masked_concat(keys, outkey='text', **masks):
+  assert all(len(keys) == len(m) for m in masks.values()), (keys, masks)
+  def _masked_concat(data):
+    # Refer to original inputs to support using same key as input/output.
+    inputs = dict(**data)
+    data[outkey] = tf.concat([inputs[k] for k in keys], axis=0)
+    for mask_name, mask_vals in masks.items():
+      m = [tf.fill(tf.shape(inputs[k]), v) for k, v in zip(keys, mask_vals)]
+      data[mask_name] = tf.concat(m, axis=0)
+    return data
+  return _masked_concat
+
+
+@Registry.register('preprocess_ops.strjoin')
+@utils.InKeyOutKey()
+def get_strjoin(glue, axis=None):
+  def _strjoin(x):
+    return tf.strings.reduce_join(x, axis=axis, separator=glue)
+  return _strjoin
+
+
+@Registry.register('preprocess_ops.majority')
+@utils.InKeyOutKey()
+def get_majority():
+  def _majority(x):
+    val, _, count = tf.unique_with_counts(x)  # Sadly, stablesorted.
+    return val[tf.argmax(count)]
+  return _majority
+
+
+@Registry.register('preprocess_ops.getidx')
+def getidx(inkey, index_key, outkey=None):
+  """Indexes a tensor and stores result in outkey."""
+  def _getidx(data):
+    idx = data[index_key]
+    array = data[inkey]
+    data[outkey or inkey] = array[idx]
+    return data
+  return _getidx

Tipsomaly/model/big_vision/pp/proj/paligemma/robustness.py

@@ -0,0 +1,72 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""pp ops."""
+
+import math
+
+from big_vision.pp import utils
+from big_vision.pp.registry import Registry
+import tensorflow as tf
+
+
+@Registry.register("preprocess_ops.resize_r")
+@utils.InKeyOutKey()
+def get_resize_r(size):
+  """Like standard `resize` but randomize some of its parameters."""
+  size = utils.maybe_repeat(size, 2)
+
+  # Sadly TF won't let us pass symbolic arguments, so we need to pre-create all
+  # variants of function calls we'd like to randomize over...
+  resize_fns = [
+      lambda x, m=m, a=a: tf.image.resize(x, size, method=m, antialias=a)
+      for m in ["bilinear", "bicubic", "lanczos3", "area", "mitchellcubic"]
+      for a in [True, False]
+  ]
+
+  def _resize_r(image):
+    """Resizes image to a given size."""
+    dtype = image.dtype
+    tf_dtype = tf.type_spec_from_value(image).dtype
+    ifn = tf.random.uniform((), 0, len(resize_fns), tf.int32)
+    image = tf.switch_case(ifn, [lambda fn=fn: fn(image) for fn in resize_fns])
+    return tf.cast(tf.clip_by_value(image, tf_dtype.min, tf_dtype.max), dtype)
+
+  return _resize_r
+
+
+@Registry.register("preprocess_ops.random_jpeg")
+@utils.InKeyOutKey()
+def get_random_jpeg(p):
+  """With probability `p`, randomly encode-decode as jpeg."""
+
+  fns = [
+      lambda x: tf.image.adjust_jpeg_quality(
+          x, dct_method="INTEGER_FAST",
+          jpeg_quality=tf.random.uniform((), 75, 96, dtype=tf.int32),
+      ),
+      lambda x: tf.image.adjust_jpeg_quality(
+          x, dct_method="INTEGER_ACCURATE",
+          jpeg_quality=tf.random.uniform((), 75, 96, dtype=tf.int32),
+      ),
+  ]
+
+  def _random_jpeg(image):
+    """Resizes image to a given size."""
+    funcs = [lambda: image] + [lambda fn=fn: fn(image) for fn in fns]
+    logits = [math.log(prob) for prob in [1 - p] + [p / len(fns)] * len(fns)]
+    fn_idx = tf.random.categorical([logits], 1, dtype=tf.int32)[0, 0]
+    return tf.switch_case(fn_idx, funcs)
+
+  return _random_jpeg

Tipsomaly/model/big_vision/pp/proj/paligemma/sciqa_ops.py

@@ -0,0 +1,65 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""pp ops."""
+
+from big_vision.pp.registry import Registry
+import tensorflow as tf
+
+
+@Registry.register('preprocess_ops.sci_qa_choices_shuffle')
+def sci_qa_choices_shuffle(
+    choice_str_inkey='choices',
+    ans_inkey='answer',
+    indexed_choices_outkey='indexed_choices',
+    indexed_answer_outkey='indexed_answer',
+):
+  """Random shuffle the sci_qa's choice on the fly.
+
+  Args:
+    choice_str_inkey: the original choice list from
+      sciqa,e.g['apple','banana',..]
+    ans_inkey: the original answer from sciqa e.g. 1
+    indexed_choices_outkey: shuffled choice (with index suffix concat to string)
+      e.g."(A) banana, (B) apple"
+    indexed_answer_outkey: shuffled answer with abc index, e,g
+      1(original)->2(shuffled)->'B' (alphabet index)
+
+  Returns:
+  """
+  def _template(data):
+    alphabet = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+    abc_tensor = tf.constant([f'({a})' for a in alphabet])
+    abcans_tensor = tf.constant([f'{a}' for a in alphabet])
+    choices = data[choice_str_inkey]
+    indices = tf.range(len(choices))
+    # Shuffle the indices
+    shuffled_indices = tf.random.shuffle(indices)
+    # Use the shuffled indices to shuffle the tensor
+    shuffled_tensor = tf.gather(choices, shuffled_indices)
+
+    abc_tensor = tf.gather(abc_tensor, indices)
+
+    data[indexed_choices_outkey] = tf.strings.reduce_join(
+        tf.strings.join([abc_tensor, shuffled_tensor], separator=' '),
+        separator=', ',
+    )
+
+    answer_tensor = data[ans_inkey]
+    new_ans_indice = tf.where(tf.equal(shuffled_indices, answer_tensor))
+    new_ans_indice = tf.gather(abcans_tensor, new_ans_indice)
+    data[indexed_answer_outkey] = tf.strings.reduce_join(new_ans_indice)
+    return data
+
+  return _template

Tipsomaly/model/big_vision/pp/proj/paligemma/segmentation.py

@@ -0,0 +1,160 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Segmentation-related ops."""
+
+import functools
+
+from big_vision.pp import registry
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.io import gfile
+
+
+_KNOWN_MODELS = {
+    'oi': 'gs://big_vision/paligemma/vae-oid.npz',
+}
+
+
+@functools.cache
+def get_checkpoint(model):
+  with gfile.GFile(_KNOWN_MODELS.get(model, model), 'rb') as f:
+    return dict(np.load(f))
+
+
+@registry.Registry.register('preprocess_ops.refcoco_mask2str')
+def get_refcoco_mask2str(model='oi'):
+  """Returns op for tokenizing a mask."""
+
+  seg_tokens = tf.constant(['<seg%03d>' % i for i in range(128)])
+  loc_tokens = tf.constant(['<loc%04d>' % i for i in range(1024)])
+  checkpoint = get_checkpoint(model)
+
+  def refcoco_mask2str(data):
+
+    mask = data['objects/mask']
+    tf.ensure_shape(mask, [None, None, 3])  # requires choice()
+    sentence = data['objects/refs/sentence']
+    tf.ensure_shape(sentence, [])  # requires choice()
+    bbox = data['objects/bbox']
+    tf.ensure_shape(bbox, [4])  # requires choice()
+
+    h = tf.cast(tf.shape(mask)[0], tf.float32)
+    w = tf.cast(tf.shape(mask)[1], tf.float32)
+    y1 = tf.cast(tf.round(h * bbox[0]), tf.int32)
+    x1 = tf.cast(tf.round(w * bbox[1]), tf.int32)
+    y2 = tf.cast(tf.round(h * bbox[2]), tf.int32)
+    x2 = tf.cast(tf.round(w * bbox[3]), tf.int32)
+
+    assert mask.dtype == tf.uint8, mask.dtype
+    mask = tf.image.resize(
+        mask[None, y1:y2, x1:x2, :1],
+        [64, 64],
+        method='bilinear',
+        antialias=True,
+    ) / 255.0
+
+    mask_indices = encode_to_codebook_indices(checkpoint, mask)[0]
+    mask_string = tf.strings.reduce_join(tf.gather(seg_tokens, mask_indices))
+
+    binned_loc = tf.cast(tf.round(bbox * 1023), tf.int32)
+    binned_loc = tf.clip_by_value(binned_loc, 0, 1023)
+    loc_string = tf.strings.reduce_join(tf.gather(loc_tokens, binned_loc))
+
+    data['prefix'] = sentence
+    data['suffix'] = tf.strings.join([loc_string, mask_string])
+
+    return data
+
+  return refcoco_mask2str
+
+
+# Based on https://arxiv.org/abs/2301.02229.
+
+NUM_DOWNSAMPLE_LAYERS = 4
+NUM_RES_BLOCKS = 2
+
+
+def encode_to_codebook_indices(checkpoint, masks):
+  """Encode a batch of binary segmentation masks into 16 tokens each.
+
+  Based on code from https://arxiv.org/abs/2301.02229
+
+  Args:
+    checkpoint: model weights from PyTorch model.
+    masks: Must be in range `[0..1]`, and of shape `[None, 64, 64, 1]`.
+
+  Returns:
+    A tensor of shape `[None, 16]` with elements in `range(128)`.
+  """
+
+  # We require that the input masks are already resized to 64x64.
+  x = tf.ensure_shape(masks, [None, 64, 64, 1])
+  x = _norm(x)
+
+  for n in range(NUM_DOWNSAMPLE_LAYERS):
+    x = _conv_tf(
+        checkpoint, x, strides=2, padding='SAME', layer_name=f'encoder.{2*n}'
+    )
+    x = tf.nn.relu(x)
+
+  for n in range(NUM_RES_BLOCKS):
+    x = _resblock_tf(checkpoint, x, layer_name=f'encoder.{8+n}.net')
+
+  x = _conv_tf(
+      checkpoint, x, strides=1, padding='SAME', layer_name='encoder.10'
+  )
+
+  return _get_codebook_indices(checkpoint, x)
+
+
+def _norm(x):
+  return 2.0 * (x - 0.5)
+
+
+def _conv_tf(checkpoint, x, strides, padding, layer_name):
+  kernel = checkpoint[layer_name + '.weight']
+  kernel = np.transpose(kernel, (2, 3, 1, 0))
+  bias = checkpoint[layer_name + '.bias']
+  return tf.nn.conv2d(x, kernel, strides=strides, padding=padding) + bias
+
+
+def _resblock_tf(checkpoint, x, layer_name):
+  """Apply a residual block of the mask encoder."""
+  original_x = x
+  x = _conv_tf(
+      checkpoint, x, padding='SAME', strides=1, layer_name=layer_name + '.0'
+  )
+  x = tf.nn.relu(x)
+  x = _conv_tf(
+      checkpoint, x, padding='SAME', strides=1, layer_name=layer_name + '.2'
+  )
+  x = tf.nn.relu(x)
+  x = _conv_tf(
+      checkpoint, x, padding='SAME', strides=1, layer_name=layer_name + '.4'
+  )
+  return x + original_x
+
+
+def _get_codebook_indices(checkpoint, encoder_output):
+  embeddings = checkpoint['_vq_vae._embedding']
+  flat_input = tf.reshape(encoder_output, [-1, embeddings.shape[1]])
+  distances = (
+      tf.reduce_sum(flat_input**2, axis=1, keepdims=True)
+      + tf.reduce_sum(embeddings**2, axis=1)
+      - 2 * tf.matmul(flat_input, embeddings.T)
+  )
+  indices = tf.argmin(distances, axis=1)
+  return tf.reshape(indices, [-1, 16])

Tipsomaly/model/big_vision/pp/proj/paligemma/video.py

@@ -0,0 +1,103 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Preprocessing for videos."""
+
+from big_vision.pp import utils
+from big_vision.pp.registry import Registry
+
+import tensorflow as tf
+
+
+@Registry.register('preprocess_ops.video_decode')
+def video_decode(res):
+  """Preprocessing."""
+
+  def _pp_per_image(img):
+    # decode
+    return tf.image.resize(tf.io.decode_jpeg(img), (res, res))
+
+  def _pp(data):
+    images = data['episodic_images']
+    # resize
+    images = tf.map_fn(_pp_per_image, images, fn_output_signature=tf.float32)
+    # rescale
+    images = 2 * (images / 255.) - 1.0
+    data['image'] = images
+    return data
+
+  return _pp
+
+
+@Registry.register('preprocess_ops.video_ensure_shape')
+def video_ensure_shape(key, shape):
+  """Preprocessing."""
+  def _video_ensure_shape(data):
+    data[key] = tf.ensure_shape(data[key], shape)
+    return data
+
+  return _video_ensure_shape
+
+
+@Registry.register('preprocess_ops.video_replicate_img')
+def video_replicate_img(replicas, num_frames):
+  """Ensure that for short videos, we have the correct number of frames.
+
+  We replicate and select.
+
+  Args:
+    replicas: num_replicas before selection. Should be less than num_frames.
+    num_frames: number of frames
+
+  Returns:
+    _replicate_img: preprocessing function
+  """
+
+  def _replicate_img(data):
+    # visual analogies + query
+    image = data['image']
+    image = tf.tile(image, [replicas, 1, 1, 1])
+    data['image'] = image[:num_frames]
+    return data
+
+  return _replicate_img
+
+
+@Registry.register('preprocess_ops.video_choice')
+@utils.InKeyOutKey()
+def video_choice(empty_fallback=None):
+  """Randomly takes one entry out of a tensor after flattening."""
+
+  def _choice(x):
+    x = tf.reshape(x, (-1,))  # Ensure it's a 1D array
+
+    # Append the fallback value so we gracefully handle empty cases.
+    x0 = tf.zeros(1, x.dtype) if empty_fallback is None else [empty_fallback]
+    x = tf.concat([x, x0], axis=0)
+
+    num_choices = tf.maximum(tf.shape(x)[0] - 1, 1)  # Don't sample x0.
+    return x[tf.random.uniform([], 0, num_choices, dtype=tf.int32)]
+
+  return _choice
+
+
+@Registry.register('preprocess_ops.stack_images')
+def stack_images(inkeys=(), outkey='image'):
+
+  def _pp(data):
+    images = tf.stack([data[inkey] for inkey in inkeys])
+    data[outkey] = images
+    return data
+
+  return _pp

Tipsomaly/model/big_vision/pp/proj/paligemma/widgetcap.py

@@ -0,0 +1,36 @@
+# Copyright 2024 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Widgetcap pp ops."""
+
+from big_vision.pp.registry import Registry
+import tensorflow as tf
+
+
+@Registry.register("preprocess_ops.draw_bbox")
+def get_draw_bbox(image_key="image", bbox_key="bbox"):
+  """Draw a single bounding box."""
+
+  def _draw_bbox(data):
+    """Draw a single bounding box."""
+    image = tf.cast(data[image_key], tf.float32)
+    image = tf.image.draw_bounding_boxes(
+        tf.expand_dims(image, 0),
+        tf.reshape(data[bbox_key], [1, 1, 4]),
+        tf.constant([255, 0, 0], dtype=tf.float32, shape=[1, 3]),
+    )
+    data[image_key] = tf.squeeze(image)
+    return data
+
+  return _draw_bbox

Tipsomaly/model/big_vision/pp/proj/uvim/pp_ops.py

@@ -0,0 +1,206 @@
+# Copyright 2022 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Preprocessing ops."""
+from big_vision.pp import utils
+from big_vision.pp.registry import Registry
+import numpy as np
+import tensorflow as tf
+
+
+@Registry.register("preprocess_ops.rgb_to_grayscale_to_rgb")
+@utils.InKeyOutKey(indefault="image", outdefault="image")
+def get_rgb_to_grayscale_to_rgb():
+  def _rgb_to_grayscale_to_rgb(image):
+    return tf.image.grayscale_to_rgb(tf.image.rgb_to_grayscale(image))
+  return _rgb_to_grayscale_to_rgb
+
+
+@Registry.register("preprocess_ops.nyu_eval_crop")
+def get_nyu_eval_crop():
+  """Crops labels and image to valid eval area."""
+  # crop_h = slice(45, 471)
+  # crop_w = slice(41, 601)
+  crop_h_start = 54
+  crop_h_size = 426
+  crop_w_start = 41
+  crop_w_size = 560
+
+  def _pp(data):
+    tf.debugging.assert_equal(tf.shape(data["labels"]), (480, 640, 1))
+    tf.debugging.assert_equal(tf.shape(data["image"]), (480, 640, 3))
+    data["labels"] = tf.slice(data["labels"],
+                              [crop_h_start, crop_w_start, 0],
+                              [crop_h_size, crop_w_size, -1])
+    data["image"] = tf.slice(data["image"],
+                             [crop_h_start, crop_w_start, 0],
+                             [crop_h_size, crop_w_size, -1])
+    return data
+  return _pp
+
+
+@Registry.register("preprocess_ops.nyu_depth")
+@utils.InKeyOutKey(indefault="depth", outdefault="labels")
+def get_nyu_depth():
+  """Preprocesses NYU depth data."""
+  def _pp(depth):
+    return tf.expand_dims(tf.cast(depth, tf.float32), -1)
+  return _pp
+
+
+@Registry.register("preprocess_ops.coco_panoptic")
+def get_coco_panoptic_pp():
+  """COCO-panoptic: produces a mask with labels and a mask with instance ids.
+
+  Instance channel will have values between 1 and N, and -1 for non-annotated
+  pixels.
+
+  Returns:
+    COCO panoptic preprocessign op.
+  """
+  def _coco_panoptic(data):
+    instance_ids = tf.cast(data["panoptic_objects"]["id"], tf.int32)
+    instance_labels = tf.cast(data["panoptic_objects"]["label"], tf.int32)
+
+    # Convert image with ids split in 3 channels into a an integer id.
+    id_mask = tf.einsum(
+        "hwc,c->hw",
+        tf.cast(data["panoptic_image"], tf.int32),
+        tf.constant([1, 256, 256**2], tf.int32))
+
+    # Broadcast into N boolean masks one per instance_id.
+    n_masks = tf.cast(
+        id_mask[:, :, None] == instance_ids[None, None, :], tf.int32)
+
+    # Merge into a semantic and an instance id mask.
+    # Note: pixels which do not belong to any mask, will have value=-1
+    # which creates an empty one_hot masks.
+    # Number instances starting at 1 (0 is treated specially by make_canonical).
+    instance_idx = tf.range(tf.shape(instance_ids)[-1])
+    instances = tf.einsum("hwc,c->hw", n_masks, instance_idx + 1)
+    semantics = tf.einsum("hwc,c->hw", n_masks, instance_labels + 1)
+
+    data["instances"] = instances[:, :, None]
+    data["semantics"] = semantics[:, :, None]
+    return data
+
+  return _coco_panoptic
+
+
+@Registry.register("preprocess_ops.make_canonical")
+@utils.InKeyOutKey(indefault="labels", outdefault="labels")
+def get_make_canonical(random=False, main_sort_axis="y"):
+  """Makes id mask ordered from left to right based on the center of mass."""
+  # By convention, instances are in the last channel.
+  def _make_canonical(image):
+    """Op."""
+    instimg = image[..., -1]
+
+    # Compute binary instance masks. Note, we do not touch 0 and neg. ids.
+    ids = tf.unique(tf.reshape(instimg, [-1])).y
+    ids = ids[ids > 0]
+    n_masks = tf.cast(
+        instimg[None, :, :] == ids[:, None, None], tf.int32)
+
+    if not random:
+      f = lambda x: tf.reduce_mean(tf.cast(tf.where(x), tf.float32), axis=0)
+      centers = tf.map_fn(f, tf.cast(n_masks, tf.int64), dtype=tf.float32)
+      centers = tf.reshape(centers, (tf.shape(centers)[0], 2))
+      major = {"y": 0, "x": 1}[main_sort_axis]
+      perm = tf.argsort(
+          centers[:, 1 - major] +
+          tf.cast(tf.shape(instimg)[major], tf.float32) * centers[:, major])
+      n_masks = tf.gather(n_masks, perm)
+    else:
+      n_masks = tf.random.shuffle(n_masks)
+
+    idx = tf.range(tf.shape(ids)[0])
+    can_mask = tf.einsum("chw,c->hw", n_masks, idx + 2) - 1
+    # Now, all 0 and neg. ids have collapsed to -1. Thus, we recover 0 id from
+    # the original mask.
+    can_mask = tf.where(instimg == 0, 0, can_mask)
+    return tf.concat([image[..., :-1], can_mask[..., None]], axis=-1)
+
+  return _make_canonical
+
+
+@Registry.register("preprocess_ops.inception_box")
+def get_inception_box(
+    *, area=(0.05, 1.0), aspect=(0.75, 1.33), min_obj_cover=0.0,
+    outkey="box", inkey="image"):
+  """Creates an inception style bounding box which can be used to crop."""
+  def _inception_box(data):
+    _, _, box = tf.image.sample_distorted_bounding_box(
+        tf.shape(data[inkey]),
+        area_range=area,
+        aspect_ratio_range=aspect,
+        min_object_covered=min_obj_cover,
+        bounding_boxes=(data["objects"]["bbox"][None, :, :]
+                        if min_obj_cover else tf.zeros([0, 0, 4])),
+        use_image_if_no_bounding_boxes=True)
+    # bbox is [[[y0,x0,y1,x1]]]
+    data[outkey] = (box[0, 0, :2], box[0, 0, 2:] - box[0, 0, :2])
+    return data
+  return _inception_box
+
+
+@Registry.register("preprocess_ops.crop_box")
+@utils.InKeyOutKey(with_data=True)
+def get_crop_box(*, boxkey="box"):
+  """Crops an image according to bounding box in `boxkey`."""
+  def _crop_box(image, data):
+    shape = tf.shape(image)[:-1]
+    begin, size = data[boxkey]
+    begin = tf.cast(begin * tf.cast(shape, tf.float32), tf.int32)
+    size = tf.cast(size * tf.cast(shape, tf.float32), tf.int32)
+    begin = tf.concat([begin, tf.constant((0,))], axis=0)
+    size = tf.concat([size, tf.constant((-1,))], axis=0)
+    crop = tf.slice(image, begin, size)
+    # Unfortunately, the above operation loses the depth-dimension. So we need
+    # to restore it the manual way.
+    crop.set_shape([None, None, image.shape[-1]])
+    return crop
+  return _crop_box
+
+
+@Registry.register("preprocess_ops.randu")
+def get_randu(key):
+  """Creates a random uniform float [0, 1) in `key`."""
+  def _randu(data):
+    data[key] = tf.random.uniform([])
+    return data
+  return _randu
+
+
+@Registry.register("preprocess_ops.det_fliplr")
+@utils.InKeyOutKey(with_data=True)
+def get_det_fliplr(*, randkey="fliplr"):
+  """Flips an image horizontally based on `randkey`."""
+  # NOTE: we could unify this with regular flip when randkey=None.
+  def _det_fliplr(orig_image, data):
+    flip_image = tf.image.flip_left_right(orig_image)
+    flip = tf.cast(data[randkey] > 0.5, orig_image.dtype)
+    return flip_image * flip + orig_image * (1 - flip)
+  return _det_fliplr
+
+
+@Registry.register("preprocess_ops.strong_hash")
+@utils.InKeyOutKey(indefault="tfds_id", outdefault="tfds_id")
+def get_strong_hash():
+  """Preprocessing that hashes a string."""
+  def _strong_hash(string):
+    return tf.strings.to_hash_bucket_strong(
+        string,
+        np.iinfo(int).max, [3714561454027272724, 8800639020734831960])
+  return _strong_hash

Tipsomaly/model/big_vision/pp/proj/uvim/pp_ops_test.py

@@ -0,0 +1,128 @@
+# Copyright 2022 Big Vision Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for pp_ops."""
+import copy
+
+from big_vision.pp.proj.uvim import pp_ops as pp
+import numpy as np
+import tensorflow as tf
+
+
+def get_image_data(dtype=tf.uint8):
+  img = tf.random.uniform((640, 320, 3), 0, 255, tf.int32)  # Can't ask uint8!?
+  return {"image": tf.cast(img, dtype)}
+
+
+class PreprocessOpsTest(tf.test.TestCase):
+
+  def tfrun(self, ppfn, data={}):  # pylint: disable=dangerous-default-value
+    # Run once as standalone, as could happen eg in colab.
+    yield {k: np.array(v) for k, v in ppfn(copy.deepcopy(data)).items()}
+
+    if not data:  # tf.data doesn't like completely empty dict...k
+      data = {"dummy": 0.0}
+
+    # And then once again as part of tfdata pipeline.
+    # You'd be surprised how much these two differ!
+    tfdata = tf.data.Dataset.from_tensors(copy.deepcopy(data))
+    for npdata in tfdata.map(ppfn).as_numpy_iterator():
+      yield npdata
+
+  def test_randu(self):
+    for output in self.tfrun(pp.get_randu("flip")):
+      self.assertEqual(output["flip"].shape, ())
+      self.assertAllGreaterEqual(output["flip"], 0.0)
+      self.assertAllLessEqual(output["flip"], 1.0)
+
+  def test_det_flip_lr(self):
+    # Test both dtypes to make it can be applied correctly to both.
+    for dtype in [tf.uint8, tf.float32]:
+      image_data = get_image_data(dtype)
+      for out in self.tfrun(pp.get_det_fliplr(randkey="rand"),
+                            {"rand": 0.1, **image_data}):
+        self.assertTrue(np.all(image_data["image"] == out["image"]))
+        self.assertEqual(out["image"].dtype, dtype)
+      for out in self.tfrun(pp.get_det_fliplr(randkey="rand"),
+                            {"rand": 0.6, **image_data}):
+        self.assertTrue(np.all(image_data["image"][:, ::-1, :] == out["image"]))
+        self.assertEqual(out["image"].dtype, dtype)
+
+  def test_inception_box(self):
+    for out in self.tfrun(pp.get_inception_box(), get_image_data()):
+      self.assertEqual(out["box"][0].shape, (2,))
+      self.assertEqual(out["box"][1].shape, (2,))
+
+  def test_crop_box(self):
+    data = get_image_data()
+    data["box"] = (tf.constant([0.5, 0.4]), tf.constant([0.25, 0.3]))
+    for out in self.tfrun(pp.get_crop_box(), data):
+      self.assertEqual(out["image"].shape, (160, 96, 3))
+      self.assertAllEqual(
+          data["image"][320:320 + 160, 128:128 + 96],
+          out["image"])
+
+  def test_make_canonical(self):
+    orig = np.array([
+        [1, 0, 3, 3, -1],
+        [1, 0, 3, 3, -1],
+        [1, 0, 2, 2, 2],
+        [1, 0, 0, -1, -1]
+    ], np.int32)[:, :, None]
+    expected = np.array([
+        [2, 0, 1, 1, -1],
+        [2, 0, 1, 1, -1],
+        [2, 0, 3, 3, 3],
+        [2, 0, 0, -1, -1]
+    ], np.int32)[:, :, None]
+    for out in self.tfrun(pp.get_make_canonical(), {"labels": orig}):
+      self.assertTrue(np.all(out["labels"] == expected))
+
+    # Test it only affects last channel.
+    for out in self.tfrun(pp.get_make_canonical(),
+                          {"labels": tf.tile(orig, (1, 1, 3))}):
+      self.assertAllEqual(out["labels"][..., 0], orig[..., 0])
+      self.assertAllEqual(out["labels"][..., 1], orig[..., 0])
+      self.assertAllEqual(out["labels"][..., 2], expected[..., 0])
+
+  def test_nyu_depth(self):
+    image = tf.zeros((5, 7, 3), dtype=tf.uint8)
+    depth = tf.zeros((5, 7), dtype=tf.float16)
+    data = {
+        "image": image,
+        "depth": depth
+    }
+    output = pp.get_nyu_depth()(data)
+    self.assertEqual(output["image"].shape, (5, 7, 3))
+    self.assertEqual(output["image"].dtype, tf.uint8)
+    self.assertEqual(output["labels"].shape, (5, 7, 1))
+    self.assertEqual(output["labels"].dtype, tf.float32)
+
+  def test_nyu_eval_crop(self):
+    image = tf.zeros((480, 640, 3), dtype=tf.uint8)
+    depth = tf.zeros((480, 640), dtype=tf.float16)
+    data = {
+        "image": image,
+        "depth": depth
+    }
+    data = pp.get_nyu_depth()(data)
+    output = pp.get_nyu_eval_crop()(data)
+    self.assertEqual(output["image"].shape, (426, 560, 3))
+    self.assertEqual(output["image"].dtype, tf.uint8)
+    self.assertEqual(output["labels"].shape, (426, 560, 1))
+    self.assertEqual(output["labels"].dtype, tf.float32)
+
+
+if __name__ == "__main__":
+  tf.test.main()