Function test_sliding_window

gemma/gm/nn/gemma3n/_modules_test.py:322–370

()

def test_sliding_window():
  num_heads = 2
  head_dim = 4
  features = 8
  cache_size = 7
  batch_size = 2
  query_pre_attn_scalar = head_dim**-0.5
  seq_len = 6

  attn = _modules.Attention(
      num_heads=num_heads,
      num_kv_heads=num_heads,
      features=features,
      head_dim=head_dim,
      attn_type=_ATTN_TYPE,
      query_pre_attn_scalar=query_pre_attn_scalar,
  )

  rng = jax.random.PRNGKey(0)
  x = jnp.ones((batch_size, seq_len, features))
  segment_pos = (
      jnp.repeat(jnp.arange(seq_len), batch_size).reshape(seq_len, batch_size).T
  )
  cache = _modules.Attention.init_cache(
      cache_size=cache_size,
      num_heads=num_heads,
      head_dim=head_dim,
      batch_size=batch_size,
      dtype=jnp.float32,
  )
  attn_mask = jnp.ones((batch_size, seq_len, cache_size))

  params = attn.init(rng, x, segment_pos, cache, attn_mask)
  _, output = attn.apply(params, x, segment_pos, cache, attn_mask)

  sliding_attn = _modules.Attention(
      num_heads=num_heads,
      num_kv_heads=num_heads,
      features=features,
      head_dim=head_dim,
      attn_type=_modules.AttentionType.LOCAL_SLIDING,
      sliding_window_size=2,
      query_pre_attn_scalar=query_pre_attn_scalar,
  )
  _, sliding_output = sliding_attn.apply(
      params, x, segment_pos, cache, attn_mask
  )

  assert not (output == sliding_output).all()


def test_query_pre_attn_scalar_modifies_output():