def _bigbird_block_rand_mask(
        from_seq_length,
        to_seq_length,
        from_block_size,
        to_block_size,
        num_rand_blocks,
        indices_prng_key: Optional[jax.random.PRNGKey] = None,
        deterministic: Optional[bool] = True,
        last_idx: Optional[int] = -1,
    ):
        """
        Create adjacency list of random attention.

        Args:
            from_seq_length: int. length of from sequence.
            to_seq_length: int. length of to sequence.
            from_block_size: int. size of block in from sequence.
            to_block_size: int. size of block in to sequence.
            num_rand_blocks: int. Number of random chunks per row.
            indices_prng_key: jax.random.PRNGKey. PRNG key that is used to perform random jax operations.
            deterministic: bool. When False random attention will be used.
            last_idx: if -1 then num_rand_blocks blocks chosen anywhere in to sequence,
            if positive then num_rand_blocks blocks chosen only up to last_idx.

        Returns:
            adjacency list of size from_seq_length//from_block_size-2 by num_rand_blocks
        """
        # using this method when from_seq_length in [1024, 3072, 4096]

        if from_seq_length // from_block_size != to_seq_length // to_block_size:
            raise ValueError("Error the number of blocks needs to be same!")
        rand_attn = jnp.zeros((from_seq_length // from_block_size - 2, num_rand_blocks), dtype=jnp.int32)
        # deterministic nor randomness
        if deterministic:
            return rand_attn

        middle_seq = jnp.arange(1, to_seq_length // to_block_size - 1, dtype=jnp.int32)
        last = to_seq_length // to_block_size - 1
        if last_idx > (2 * to_block_size):
            last = (last_idx // to_block_size) - 1

        r = num_rand_blocks  # shorthand
        for i in range(1, from_seq_length // from_block_size - 1):
            start = i - 2
            end = i
            if i == 1:
                seq_values = jax.random.permutation(indices_prng_key, middle_seq[2:last])[:r]
                rand_attn = rand_attn.at[i - 1].set(seq_values)
            elif i == 2:
                seq_values = jax.random.permutation(indices_prng_key, middle_seq[3:last])[:r]
                rand_attn = rand_attn.at[i - 1].set(seq_values)
            elif i == from_seq_length // from_block_size - 3:
                seq_values = jax.random.permutation(indices_prng_key, middle_seq[:last])[:r]
                rand_attn = rand_attn.at[i - 1].set(seq_values)
            # Missing -3: should have been sliced till last-3
            elif i == from_seq_length // from_block_size - 2:
                seq_values = jax.random.permutation(indices_prng_key, middle_seq[:last])[:r]
                rand_attn = rand_attn.at[i - 1].set(seq_values)
            # Missing -4: should have been sliced till last-4
            else:
                if start > last:
                    start = last
                    seq_values = jax.random.permutation(indices_prng_key, middle_seq[:start])[:r]
                    rand_attn = rand_attn.at[i - 1].set(seq_values)
                elif (end + 1) == last:
                    seq_values = jax.random.permutation(indices_prng_key, middle_seq[:start])[:r]
                    rand_attn = rand_attn.at[i - 1].set(seq_values)
                else:
                    concat_values = jnp.concatenate((middle_seq[:start], middle_seq[end + 1 : last]))
                    seq_values = jax.random.permutation(indices_prng_key, concat_values)[:r]
                    rand_attn = rand_attn.at[i - 1].set(seq_values)
        return rand_attn

        Create adjacency list of random attention.

        Args:
            from_seq_length: int. length of from sequence.
            to_seq_length: int. length of to sequence.
            from_block_size: int. size of block in from sequence.
            to_block_size: int. size of block in to sequence.
            num_rand_blocks: int. Number of random chunks per row.
            indices_prng_key: jax.random.PRNGKey. PRNG key that is used to perform random jax operations.
            deterministic: bool. When False random attention will be used.
            last_idx: if -1 then num_rand_blocks blocks chosen anywhere in to sequence,
            if positive then num_rand_blocks blocks chosen only up to last_idx.

        Returns:
            adjacency list of size from_seq_length//from_block_size-2 by num_rand_blocks
        
dpkd/transformers/src/transformers/models/big_bird/modeling_flax_big_bird.py