# Copyright 2023 The JAX 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
#
#     https://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.

"""Mini-mask creation library."""

from __future__ import annotations

import collections
import functools
from typing import Callable, Dict, List, NamedTuple, Set, Tuple
from jax import util as jax_util
from fjformer.pallas_operations.tpu.splash_attention import (
	splash_attention_mask as mask_lib,
)
import numpy as np


# mypy: ignore-errors


# Logic for processing NumPy masks for kernels
class MaskInfo(NamedTuple):
	"""Contains runtime masking information for the Splash attention kernel.

	The arrays data_next, mask_next and block_mask are placed in TPU
	scalar-memory. This is a scarse resource so the mask creation logic attempts
	to shrink the data-type of these arrays to the smallest possible one.
	This can be: np.int32, np.int16 or np.int8.

	For the arrays data_next, mask_next and block_mask the size of the first
	dimension can be one of the two following values: num_head or
	num_head_shards.
	The first dimension has size:
	* num_head_shards when there is only one unique mask for each head in a shard.
	In this case the three arrays are broadcasted to all the heads in the shard.
	* num_heads when there is more than one unique mask for each head in the
	shard.

	Attributes:
	  data_next: An integer[num_heads_or_shards, num_q_blocks, num_kv_blocks]
	    NumPy array where each entry contains the next `kv` block index to
	    prefetch.
	  mask_next: An integer[num_heads_or_shards, num_q_blocks, num_kv_blocks]
	    NumPy array where each entry contains the next mask block index in
	    `partial_mask_blocks` to prefetch.
	  block_mask: An integer[num_heads_or_shards, num_q_blocks, num_kv_blocks]
	    NumPy array whose entries can be 0, 1 or 2. An entry of 0 indicates that
	    the corresponding block in the full mask was all zeros. An entry of 1
	    indicates that the corresponding block in the full mask contained both
	    zeros and ones. An entry of 2 indicates the corresponding block was
	    entirely ones.
	  partial_mask_blocks: A i32[num_partial_blocks, block_q, block_kv] NumPy
	    array that contains the blocks of the original mask that contained both
	    zeros and ones. The entries in `mask_next` point to indices in the first
	    axis of this array.
	  q_sequence: A i32[q_sequence_length] NumPy array. When using causal masking,
	    this contains the list of indices that correspond to q tokens. For plain
	    causal this is just np.arange(q_sequence_length).
	"""

	data_next: np.ndarray | None
	mask_next: np.ndarray | None
	block_mask: np.ndarray | None
	partial_mask_blocks: np.ndarray | None
	q_sequence: np.ndarray | None


def _downcast_to_small_type(array: np.ndarray) -> np.ndarray:
	"""Downcast numpy array.

	If possible, downcast the data-type of the input array to the smallest numpy
	type (among np.int16 and np.int8) that fits the content of the array.

	Args:
	  array: the array to downcast

	Returns:
	  The downcasted array.

	Raises:
	  ValueError: if the input array is not np.int32 or if its elements are not
	  all positive.
	"""

	if array.dtype != np.int32:
		raise ValueError("Expected int32 input.")

	if not np.all(array >= 0):
		raise ValueError("Expected non-negative array.")

	if array.size == 0:
		return array

	max_value = np.max(array)

	if max_value <= np.iinfo(np.int8).max:
		return array.astype(np.int8)
	elif max_value <= np.iinfo(np.int16).max:
		return array.astype(np.int16)
	else:
		return array.astype(np.int32)


def _check_mask(mask: mask_lib.Mask) -> None:
	"""Check that the given mask is valid.

	A row of all zeros along the kv dimension would result in a division by zero
	when computing the softmax. This function is meant to protect against that
	case.

	Args:
	  mask: the mask to check.

	Raises:
	  ValueError: the mask is invalid.
	"""

	assert len(mask.shape) == 2

	exception_message = (
		"Some rows of the mask (along the kv dimension) are all zeros.\nThis is"
		" would result in a division by zero when computing the attention"
		" softmax."
	)

	is_row_non_zero = np.zeros(mask.shape[0], dtype=np.bool_)
	for col in range(mask.shape[1]):
		# Mask only supports slice indices.
		is_row_non_zero = np.logical_or(
			is_row_non_zero,
			mask[(slice(0, mask.shape[0]), slice(col, col + 1))][:, 0],
		)
	if not is_row_non_zero.all():
		raise ValueError(exception_message)


class _HashableNDArray:
	"""Helper to make a numpy array hashable: can be added associative containers.

	Attributes:
	  array: The underlying numpy array.
	"""

	array: np.ndarray

	def __init__(self, array: np.ndarray):
		self.array = array

	def __hash__(self):
		return hash(self.array.tobytes())

	def __eq__(self, other: object) -> bool:
		if not isinstance(other, _HashableNDArray):
			return NotImplemented
		return np.array_equal(self.array, other.array, equal_nan=True)


def _get_mask_info_for_shard(
	output_shape: Tuple[int, int, int],
	has_mask_next: bool,
	mask: mask_lib.MultiHeadMask,
	block_shape: Tuple[int, int],
	coords_to_partial_mask_block_index: Dict[Tuple[int, int, int], int],
	masks_per_head_shard: int,
	head_start: int,
	num_heads: int,
	q_seq_start: int,
	q_seq_shard_size: int,
	blocked_q_seq_start: int,
	is_dkv: bool,
) -> Tuple[np.ndarray, np.ndarray | None]:
	"""Process a slice of the mask to compute data_next and mask_next.

	Args:
	  output_shape: The shape of the data_next and mask_next to return
	  has_mask_next: Whether mask_next should be constructed. If False None is
	    returned for mask_next.
	  mask: The full mask to be sliced according to the head and sequence ranges
	  block_shape: Shape of the Pallas grid block.
	  coords_to_partial_mask_block_index: Mapping between the pallas launch grid
	    coordinates and the index of the corresponding block in partial mask block
	    list.
	  masks_per_head_shard: Number of masks per head shards
	  head_start: First head of the current shard.
	  num_heads: Number of heads in the shard.
	  q_seq_start: Start index along the Q sequence for the current shard (in
	    number of tokens).
	  q_seq_shard_size: Number of tokens along the Q sequence for the current
	    shard.
	  blocked_q_seq_start: Start index along the Q sequence for the current shard
	    (in number of grid blocks)
	  is_dkv: True if we are processing the dKV mask

	Returns:
	  Slice of data_next and mask_next (if required) that correspond to the
	  current mask slice.
	"""
	_, _, kv_seq_len = mask.shape

	q_block_size, kv_block_size = block_shape
	q_block_count, q_mod = divmod(q_seq_shard_size, q_block_size)
	kv_block_count, kv_mod = divmod(kv_seq_len, kv_block_size)

	assert q_mod == 0
	assert kv_mod == 0

	blocked_shape = (
		(kv_block_count, num_heads, q_block_count)
		if is_dkv
		else (num_heads, q_block_count, kv_block_count)
	)

	data_coords = []
	mask_coords = []
	for idx in np.ndindex(blocked_shape):
		if is_dkv:
			kv_index, h_index, q_index = idx
		else:
			h_index, q_index, kv_index = idx

		h_index = h_index if masks_per_head_shard == 1 else head_start + h_index

		chunk = mask[
			(
				h_index,
				slice(
					q_seq_start + q_index * q_block_size,
					q_seq_start + (q_index + 1) * q_block_size,
				),
				slice(kv_index * kv_block_size, (kv_index + 1) * kv_block_size),
			)
		]
		if chunk.any():
			data_coords.append(idx)
			if not chunk.all():
				mask_coords.append(idx)

	# Initialize the output arrays.
	mask_next = None
	if has_mask_next:
		mask_next = np.zeros(output_shape, dtype=np.int32)
	data_next = np.zeros(output_shape, dtype=np.int32)

	# If the mask is completelly zero'd out return freshly initialized outputs.
	if not data_coords:
		return data_next, mask_next

	data_coords_iter = iter(data_coords)
	first_j = coord_j = next(data_coords_iter)
	if mask_next is not None and mask_coords:
		mask_coords_iter = iter(mask_coords)
		first_m = coord_m = next(mask_coords_iter)
	else:
		first_m, coord_m, mask_coords_iter = None, None, None

	for idx in np.ndindex(blocked_shape):
		if is_dkv:
			kv_index, h_index, q_index = idx
			chunk_idx: tuple[int, ...] = (h_index, q_index, kv_index)
			data_dim = 2
		else:
			chunk_idx = idx
			data_dim = 2

		is_next = idx > coord_j

		if is_next:
			try:
				coord_j = next(data_coords_iter)
			except StopIteration:
				coord_j = first_j
		data_next[chunk_idx] = coord_j[data_dim]

		if mask_next is not None and mask_coords:
			assert coord_m is not None
			is_next_mask = idx > coord_m
			if is_next_mask:
				try:
					coord_m = next(mask_coords_iter)  # type: ignore
				except StopIteration:
					coord_m = first_m

			if is_dkv:
				assert coord_m is not None
				coord_m_global = (
					coord_m[1] + head_start,
					coord_m[2] + blocked_q_seq_start,
					coord_m[0],
				)
			else:
				assert coord_m is not None
				coord_m_global = (
					coord_m[0] + head_start,
					coord_m[1] + blocked_q_seq_start,
					coord_m[2],
				)

			mask_next[chunk_idx] = coords_to_partial_mask_block_index[coord_m_global]

	return data_next, mask_next


# When used in a transformer network with multiple layers, the SplashAttention
# kernel is created serveral times with the same mask. Cache MaskInfo to avoid
# blowing up compile times. Ideally the size of the cache should be determined
# by the client.
@functools.lru_cache(maxsize=12)
def _process_mask(
	mask: mask_lib.MultiHeadMask,  # [num_heads, q_seq_len, kv_seq_len]
	block_shape: Tuple[int, int],
	is_dkv: bool,
	*,
	downcast_smem_data: bool = True,
	head_shards: int = 1,
	q_seq_shards: int = 1,
	shrink_grid: bool = True,
) -> tuple[MaskInfo, jax_util.HashableFunction | None]:
	"""Transform a dense mask into a sparse representation.

	The number of head and Q sequence shards are needed to create a MaskInfo
	object that is partitionable (with shmap or PartIR) along these two dimension.
	In particular for dKV MaskInfo, for each shard the indices of in the data_next
	array are relative to the current shard.
	The fwd and dQ MaskInfo objects do not change when sharding along the head or
	Q dimensions, they would be different if we were to shard along the KV
	dimension, but the kernel does not support that.

	Args:
	  mask: Dense mask to process.
	  block_shape: Shape of the Pallas grid block.
	  is_dkv: True if we are processing the dKV mask
	  downcast_smem_data: If True, downcast the scalar-memory data of MaskInfo to
	    a data type smaller than np.int32 (if possible).
	  head_shards: Number of head shards of the mesh in which the kernel is
	    launched.
	  q_seq_shards: Number of Q sequence shards of the mesh in which the kernel is
	    launched.
	  shrink_grid: Whether or not we should apply the grid shrinking optimization.

	Returns:
	  `MaskInfo`, a sparse representation of the dense mask.
	  `MaskCallable`: a callable that, given in input Q and KV indices, returns
	    the value of the mask at those coordinates.

	Raises:
	  ValueError: if the input mask is invalid or the block sizes are not
	  compatible with the mask sizes.
	"""

	if len(mask.shape) != 3:
		raise ValueError(f"Expected a 3-dim mask, instead got: {mask.shape=}")

	head_count, q_seq_len, kv_seq_len = mask.shape
	q_block_size, kv_block_size = block_shape
	q_blocks_count, q_mod = divmod(q_seq_len, q_block_size)
	kv_blocks_count, kv_mod = divmod(kv_seq_len, kv_block_size)

	if q_mod != 0:
		raise ValueError(f"{q_block_size=} should divide {q_seq_len=}.")
	if kv_mod != 0:
		raise ValueError(f"{kv_block_size=} should divide {kv_seq_len=}.")

	q_seq_len_per_shard, mod = divmod(q_seq_len, q_seq_shards)
	if mod != 0:
		raise ValueError(f"{q_seq_shards=} should divide {q_seq_len=}.")

	q_blocks_per_shard, mod = divmod(q_seq_len_per_shard, q_block_size)
	if mod != 0:
		raise ValueError(f"{q_block_size=} should divide {q_seq_len_per_shard=}.")

	heads_per_shard, mod = divmod(head_count, head_shards)
	if mod != 0:
		raise ValueError(f"{head_shards=} should divide {head_count=}.")

	first_mask_size = mask.masks[0].shape
	for h in range(head_count):
		if mask.masks[h].shape != first_mask_size:
			raise ValueError(
				f"First head mask has shape {first_mask_size}, but head mask {h} has"
				f" shape {mask.masks[h].shape}. All head masks must have the same"
				" shape."
			)

	# Uniquify the masks.
	# Create a collection of the unique head masks in the input multi-head mask.
	# This avoids processing the same mask multiple times and it enables
	# deduplicating the partial mask blocks information.
	# Assign a unique ID to every unique mask.

	def assign_unique_ids(objects):
		id_map = collections.defaultdict(lambda: len(id_map))
		return {obj: id_map[obj] for obj in objects}

	unique_masks_dict: Dict[mask_lib.Mask, int] = assign_unique_ids(
		head_mask for head_mask in mask.masks
	)

	# Build a mapping of heads to unique masks and masks to unique masks.

	head_to_mask_id: List[int] = [0] * head_count
	head_shard_to_mask_ids: List[Set[int]] = [set() for _ in range(head_shards)]
	mask_id_to_heads: List[List[int]] = [[] for _ in range(len(unique_masks_dict))]
	mask_id_to_head_shards: List[Set[int]] = [
		set() for _ in range(len(unique_masks_dict))
	]

	for head in range(head_count):
		mask_id = unique_masks_dict[mask.masks[head]]
		head_to_mask_id[head] = mask_id
		head_shard = head // heads_per_shard
		head_shard_to_mask_ids[head_shard].add(mask_id)
		mask_id_to_heads[mask_id].append(head)
		mask_id_to_head_shards[mask_id].add(head_shard)

	# If we have at most one unique mask per each head shard, then we can brodcast
	# the mask to all the heads in the shard. This is the common case.
	# If we have more than one mask in each head shard, then the optimization
	# cannot kick in and we use one mask for each head.
	# heads to masks. This will likely introduce an additional field in the
	# MaskInfo class and runtime overhead to perform an indirect lookup. Since
	# having multiple masks per head-shard is not a common case we leave this for
	# future work.
	max_masks_per_head_shard = max([len(x) for x in head_shard_to_mask_ids])
	masks_per_head_shard = 1 if max_masks_per_head_shard == 1 else heads_per_shard

	unique_masks = [
		pair[0] for pair in sorted(unique_masks_dict.items(), key=lambda x: x[1])
	]

	partial_mask_block_ids: Dict[_HashableNDArray, int] = collections.defaultdict(
		lambda: len(partial_mask_block_ids)
	)
	block_id_to_block_coords: Dict[int, List[Tuple[int, ...]]] = collections.defaultdict(
		list
	)

	# Shape of the block_mask, mask_next and data_next fields of MaskInfo.
	block_mask_shape = (
		head_shards if masks_per_head_shard == 1 else head_count,
		q_blocks_count,
		kv_blocks_count,
	)

	block_mask = np.zeros(block_mask_shape, dtype=np.int32)

	def set_block_mask(mask_id: int, q_index: int, kv_index: int, value: int):
		if masks_per_head_shard == 1:
			for shard_index in mask_id_to_head_shards[mask_id]:
				block_mask[shard_index, q_index, kv_index] = value
		else:
			for head_index in mask_id_to_heads[mask_id]:
				block_mask[head_index, q_index, kv_index] = value

	q_sequence = None
	mask_function = None
	if len(unique_masks) == 1:
		unique_mask = unique_masks[0]
		# The mask object either define q_sequence and mask_function or none of
		# them.
		assert hasattr(unique_mask, "q_sequence") == hasattr(unique_mask, "mask_function")

		# If the mask object defines a q_sequence and a mask_function, then make use
		# of these in the kernel rather. This is preferable over loading the mask
		# from memory. When using a mask_function, then mask_next and
		# partial_mask_blocks are left undefined and not used in the kernel.
		if hasattr(unique_mask, "q_sequence") and hasattr(unique_mask, "mask_function"):
			q_sequence = unique_mask.q_sequence
			mask_function = unique_mask.mask_function

	# Identify the partial mask blocks and the value of the block mask for each
	# block.
	# The partial mask blocks are "global", meaning it is a collection of such
	# blocks across all the heads. Partial mask blocks are uniquified.
	# When partitioning, all partial mask blocks are replicated across shards.
	# As a future extension we could assign to each head shard only the relevant
	# partial blocks. This would reduce the amount of data each shard needs
	# from: num_unique_partial_block_blocks to
	# max(unique_partial_blocks_per_shard)
	for mask_id, unique_mask in enumerate(unique_masks):
		for coords in np.ndindex((q_blocks_count, kv_blocks_count)):
			(q_index, kv_index) = coords
			chunk = unique_mask[
				(
					slice(q_index * q_block_size, (q_index + 1) * q_block_size),
					slice(kv_index * kv_block_size, (kv_index + 1) * kv_block_size),
				)
			]
			has_nonzero = chunk.any()
			if has_nonzero:
				all_nonzero = chunk.all()
				if not all_nonzero:
					set_block_mask(mask_id, q_index, kv_index, 1)
					partial_mask_block_id = partial_mask_block_ids[_HashableNDArray(chunk)]
					for head_index in mask_id_to_heads[mask_id]:
						block_id_to_block_coords[partial_mask_block_id].append(
							(head_index,) + coords
						)
				else:
					set_block_mask(mask_id, q_index, kv_index, 2)

	unique_partial_mask_blocks = [
		pair[0] for pair in sorted(partial_mask_block_ids.items(), key=lambda x: x[1])
	]

	# For each position in the in q, kv grid that contains a partial mask block
	# record the index of its partial mask block.
	coords_to_partial_mask_block_index = {}
	for partial_mask_block_id, coords in block_id_to_block_coords.items():
		for coo in coords:
			coords_to_partial_mask_block_index[coo] = partial_mask_block_id

	partial_mask_blocks = None
	has_mask_next = False
	if len(unique_partial_mask_blocks) == 1:
		partial_mask_blocks = [x.array for x in unique_partial_mask_blocks]
		partial_mask_blocks = partial_mask_blocks[0][None].astype(np.int32)
		has_mask_next = True
	elif len(unique_partial_mask_blocks) > 1:
		partial_mask_blocks = [x.array for x in unique_partial_mask_blocks]
		partial_mask_blocks = np.stack(partial_mask_blocks, axis=0).astype(np.int32)
		has_mask_next = True
	if is_dkv and partial_mask_blocks is not None:
		partial_mask_blocks = np.swapaxes(partial_mask_blocks, -1, -2)

	all_head_shards_identical = all(
		head_shard_to_mask_ids[0] == x and len(x) == 1 for x in head_shard_to_mask_ids
	)
	# When all the head shards are identical, we can process only one and
	# broadcast the result across all the head shards, this avoids redundant work.
	shards_to_process = 1 if all_head_shards_identical else head_shards

	# Iterate over the shards.
	# Work on a fraction of the mask at the time to compute the mask and the data
	# indices. This is needed to compute the correct data indices, which are
	# relative to the current slice of the mask.

	q_sequence_axis = 1
	head_axis = 0

	# Collect mask_info shards along the head dimension, concatentate (or
	# broadcast) them after the loop.
	data_next_per_head_list, mask_next_per_head_list = [], []
	for head_shard in range(shards_to_process):
		# Collect mask_info shards along the q_sequence dimension, concatenate them
		# after the loop.
		data_next_sequence_slices, mask_next_sequence_slices = [], []
		for q_seq_len_shard in range(q_seq_shards):
			head_start = head_shard * heads_per_shard
			q_seq_len_shard_size = q_blocks_per_shard * q_block_size
			q_seq_len_start = q_seq_len_shard * q_seq_len_shard_size

			blocked_q_seq_len_start = q_seq_len_shard * q_blocks_per_shard
			blocked_q_seq_len_slice = slice(
				blocked_q_seq_len_start,
				(q_seq_len_shard + 1) * q_blocks_per_shard,
			)

			if masks_per_head_shard == 1:
				unique_mask = unique_masks[head_to_mask_id[head_start]]
				unique_mask = mask_lib.MultiHeadMask((unique_mask,))
				current_mask = unique_mask
				mask_head_slice = slice(head_shard, head_shard + 1)
			else:
				current_mask = mask
				mask_head_slice = slice(head_start, (head_shard + 1) * heads_per_shard)

			# The current iteration of the loop processes a slice of the mask info
			# tensors of this shape:
			mask_info_slice_shape = (
				mask_head_slice.stop - mask_head_slice.start,
				blocked_q_seq_len_slice.stop - blocked_q_seq_len_slice.start,
				kv_blocks_count,
			)
			# Generate data_next and mask_next for the current slice.
			data_next_slice, mask_next_slice = _get_mask_info_for_shard(
				output_shape=mask_info_slice_shape,
				has_mask_next=has_mask_next,
				mask=current_mask,
				block_shape=block_shape,
				coords_to_partial_mask_block_index=coords_to_partial_mask_block_index,
				head_start=head_start,
				masks_per_head_shard=masks_per_head_shard,
				num_heads=1 if masks_per_head_shard == 1 else heads_per_shard,
				q_seq_start=q_seq_len_start,
				q_seq_shard_size=q_seq_len_shard_size,
				blocked_q_seq_start=blocked_q_seq_len_start,
				is_dkv=is_dkv,
			)
			data_next_sequence_slices.append(data_next_slice)
			mask_next_sequence_slices.append(mask_next_slice)

		# Concatenate the sequence shards.
		data_next_per_head = np.concatenate(data_next_sequence_slices, axis=q_sequence_axis)
		data_next_per_head_list.append(data_next_per_head)
		if has_mask_next:
			mask_next_per_head = np.concatenate(
				mask_next_sequence_slices, axis=q_sequence_axis
			)
			mask_next_per_head_list.append(mask_next_per_head)

	# Concatenate (or broadcast) the head shards.
	mask_next = None
	if all_head_shards_identical:
		assert len(data_next_per_head_list) == 1
		data_next_shard = data_next_per_head_list[0]
		assert data_next_shard.shape == (1, q_blocks_count, kv_blocks_count)
		data_next = np.broadcast_to(
			data_next_shard,
			(head_shards, q_blocks_count, kv_blocks_count),
		)
		if has_mask_next:
			assert len(mask_next_per_head_list) == 1
			mask_next_shard = mask_next_per_head_list[0]
			assert mask_next_shard.shape == (1, q_blocks_count, kv_blocks_count)
			mask_next = np.broadcast_to(
				mask_next_shard,
				(head_shards, q_blocks_count, kv_blocks_count),
			)
	else:
		data_next = np.concatenate(data_next_per_head_list, axis=head_axis)
		if has_mask_next:
			mask_next = np.concatenate(mask_next_per_head_list, axis=head_axis)

	# Shrink the width of the mask info when possible.
	if shrink_grid and block_mask.shape[0] == head_shards and len(unique_masks) == 1:
		rows_per_q_shard = block_mask.shape[1] // q_seq_shards
		block_mask_shards = []
		data_next_shards = []
		mask_next_shards = []

		# Slice the mask info arrays along the Q dimension. This simulates Q
		# sharding and enabled shrinking each Q shard of the MaskInfo arrays
		# independently.
		for q_seq_len_shard in range(q_seq_shards):
			rows = slice(
				q_seq_len_shard * rows_per_q_shard,
				(q_seq_len_shard + 1) * rows_per_q_shard,
			)
			current_block_mask = block_mask[:, rows, :]
			current_data_next = data_next[:, rows, :]
			current_mask_next = mask_next[:, rows, :] if mask_next is not None else None

			shrink_function = _shrink_mask_info_dkv if is_dkv else _shrink_mask_info

			current_block_mask, current_data_next, current_mask_next = shrink_function(
				block_mask=current_block_mask,
				data_next=current_data_next,
				mask_next=current_mask_next,
				head_shards=head_shards,
			)

			assert current_block_mask.size > 0
			assert current_data_next.size > 0
			assert current_mask_next is None or current_mask_next.size > 0
			assert current_block_mask.shape == current_data_next.shape
			assert (
				current_mask_next is None or current_block_mask.shape == current_mask_next.shape
			)

			block_mask_shards.append(current_block_mask)
			data_next_shards.append(current_data_next)
			mask_next_shards.append(current_mask_next)

		if q_seq_shards == 1:
			block_mask = block_mask_shards[0]
			data_next = data_next_shards[0]
			mask_next = mask_next_shards[0]
		else:
			# Since shrinking happens independently for each shard along Q we might
			# end up with uneven mask info shards. Insert padding where necessary to
			# maintain the SPMD paradigm.
			padding_axis = 1 if is_dkv else 2

			max_size = max([x.shape[padding_axis] for x in block_mask_shards])
			padded_block_mask_shards = []
			padded_data_next_shards = []
			padded_mask_next_shards = []
			assert len(block_mask_shards) == len(data_next_shards) == len(mask_next_shards)
			for (
				current_block_mask,
				current_data_next,
				current_mask_next,
			) in zip(block_mask_shards, data_next_shards, mask_next_shards):
				# For dKV shrinking happens along axis Q (the rows of MaskInfo), for
				# fwd and dQ shrinking happens along axis KV (the columns of MaskInfo).
				if is_dkv:
					pad_width = (
						(0, 0),
						(0, max_size - current_block_mask.shape[padding_axis]),
						(0, 0),
					)
				else:
					pad_width = (
						(0, 0),
						(0, 0),
						(0, max_size - current_block_mask.shape[padding_axis]),
					)

				padded_block_mask_shards.append(
					np.pad(current_block_mask, pad_width=pad_width, constant_values=0)
				)
				# By padding data_next and mask_next with 'edge' policy we avoid
				# fetching new blocks thanks to 'immediate revisiting'.
				padded_data_next_shards.append(
					np.pad(current_data_next, pad_width=pad_width, mode="edge")
				)
				if current_mask_next is not None:
					padded_mask_next_shards.append(
						np.pad(current_mask_next, pad_width=pad_width, mode="edge")
					)

			block_mask = np.concatenate(padded_block_mask_shards, axis=1)
			data_next = np.concatenate(padded_data_next_shards, axis=1)
			mask_next = np.concatenate(padded_mask_next_shards, axis=1)

	if downcast_smem_data:
		data_next = _downcast_to_small_type(data_next)
		block_mask = _downcast_to_small_type(block_mask)
		if mask_next is not None:
			mask_next = _downcast_to_small_type(mask_next)

	assert (mask_function is not None) == (q_sequence is not None)
	# When the mask can be computed inside the kernel with a mask_function,
	# there is no need to load it from memory. So mask_next and
	# partial_mask_blocks are unused.
	return (
		MaskInfo(
			data_next=data_next,
			mask_next=mask_next if mask_function is None else None,
			block_mask=block_mask,
			partial_mask_blocks=partial_mask_blocks if mask_function is None else None,
			q_sequence=q_sequence,
		),
		mask_function,
	)


def _shrink_mask_info(
	*,
	block_mask: np.ndarray,
	data_next: np.ndarray,
	mask_next: np.ndarray,
	head_shards: int,
):
	assert block_mask.ndim == 3
	assert data_next.ndim == 3
	assert mask_next is None or mask_next.ndim == 3

	head_block_mask = block_mask[0]

	grouped_non_zero_cols = []
	# Group non-zero columns based on which row they belong to.
	for row_index in range(head_block_mask.shape[0]):
		head_block_mask_row = head_block_mask[row_index, :]
		non_zero_cols = np.nonzero(head_block_mask_row)[0]
		grouped_non_zero_cols.append(non_zero_cols)

	# Pad each row in the non-zero indices to match the width of the longest
	# row. This avoids having jagged rows.
	max_non_zero_cols = max([len(x) for x in grouped_non_zero_cols])
	padded_non_zero_cols = []
	padding = -1
	for row in grouped_non_zero_cols:
		padded_non_zero_cols.append(
			np.pad(
				row,
				pad_width=(0, max_non_zero_cols - row.shape[0]),
				constant_values=padding,
			)
		)

	padded_non_zero_cols = np.stack(padded_non_zero_cols, axis=0)

	assert padded_non_zero_cols.shape[0] == block_mask.shape[1], (
		padded_non_zero_cols.shape,
		block_mask.shape,
	)

	# For each row of array, select the columns indices in padded_non_zero_cols,
	# ignore padding.
	def select_cols(array):
		assert array.ndim == 2
		assert padded_non_zero_cols.ndim == 2
		assert array.shape[0] == padded_non_zero_cols.shape[0]
		assert array.shape[1] >= padded_non_zero_cols.shape[1]
		selected_rows = []
		for row in range(array.shape[0]):
			col = padded_non_zero_cols[row]
			selected = array[row][col]
			selected = np.where(col != padding, selected, 0)
			selected_rows.append(selected)

		return np.stack(selected_rows, axis=0)

	return _slice_mask_info(
		block_mask=block_mask,
		data_next=data_next,
		mask_next=mask_next,
		head_shards=head_shards,
		slice_function=select_cols,
	)


def _shrink_mask_info_dkv(
	*,
	block_mask: np.ndarray,
	data_next: np.ndarray,
	mask_next: np.ndarray,
	head_shards: int,
):
	assert block_mask.ndim == 3
	assert data_next.ndim == 3
	assert mask_next is None or mask_next.ndim == 3

	head_block_mask = block_mask[0]
	grouped_non_zero_rows = []
	# Group non-zero rows based on which column they belong to.
	for col_index in range(head_block_mask.shape[1]):
		col = head_block_mask[:, col_index]
		non_zero_rows = np.nonzero(col)[0]
		grouped_non_zero_rows.append(non_zero_rows)

	# Pad each col in the non-zero indices to match the height of the longest
	# col. This avoids having jagged cols.
	max_non_zero_rows = max([len(x) for x in grouped_non_zero_rows])
	padded_non_zero_rows = []
	padding = -1
	for col in grouped_non_zero_rows:
		padded_non_zero_rows.append(
			np.pad(
				col,
				pad_width=(max_non_zero_rows - col.shape[0], 0),
				constant_values=padding,
			)
		)

	padded_non_zero_rows = np.stack(padded_non_zero_rows, axis=1)

	assert padded_non_zero_rows.shape[1] == block_mask.shape[2], (
		padded_non_zero_rows.shape,
		block_mask.shape,
	)

	# For each col of array, select the rows indices in padded_non_zero_rows,
	# ignore padding.
	def select_rows(array):
		assert array.ndim == 2
		assert padded_non_zero_rows.ndim == 2
		assert array.shape[1] == padded_non_zero_rows.shape[1]
		assert array.shape[0] >= padded_non_zero_rows.shape[0]
		selected_cols = []
		for col in range(array.shape[1]):
			row = padded_non_zero_rows[:, col]
			selected = array[:, col][row]
			selected = np.where(row != padding, selected, 0)
			selected_cols.append(selected)

		return np.stack(selected_cols, axis=1)

	return _slice_mask_info(
		block_mask=block_mask,
		data_next=data_next,
		mask_next=mask_next,
		head_shards=head_shards,
		slice_function=select_rows,
	)


def _slice_mask_info(
	*,
	block_mask: np.ndarray,
	data_next: np.ndarray,
	mask_next: np.ndarray,
	head_shards: int,
	slice_function: Callable[[np.ndarray], np.ndarray],
):
	new_block_mask = []
	new_data_next = []
	new_mask_next = []
	for head_shard in range(head_shards):
		head_block_mask = block_mask[head_shard]
		head_block_mask = slice_function(head_block_mask)
		new_block_mask.append(head_block_mask)

		head_data_next = data_next[head_shard]
		head_data_next = slice_function(head_data_next)
		new_data_next.append(head_data_next)

		if mask_next is not None:
			head_mask_next = mask_next[head_shard]
			head_mask_next = slice_function(head_mask_next)
			new_mask_next.append(head_mask_next)

	block_mask = np.stack(new_block_mask, axis=0)
	data_next = np.stack(new_data_next, axis=0)
	if mask_next is not None:
		mask_next = np.stack(new_mask_next, axis=0)

	return block_mask, data_next, mask_next


process_mask = functools.partial(_process_mask, is_dkv=False)
process_mask_dkv = functools.partial(_process_mask, is_dkv=True)
