Upload directory

models/vit_kprpe/RPE/KPRPE/kprpe_shared.py

@@ -0,0 +1,735 @@
+from easydict import EasyDict as edict
+import math
+import torch
+import torch.nn as nn
+from .dist import _rp_2d_cross_cols, _rp_2d_cross_rows, _rp_2d_euclidean, _rp_2d_product, _rp_2d_quant
+
+try:
+ from ..rpe_ops.rpe_index import RPEIndexFunction
+except Exception as e:
+ print('Failed to import cuda/cpp RPEIndexFunction')
+ RPEIndexFunction = None
+
+
+
+def get_absolute_positions(height, width, dtype, device):
+ '''Get absolute positions
+
+ Take height = 3, width = 3 as an example:
+ rows: cols:
+ 1 1 1 1 2 3
+ 2 2 2 1 2 3
+ 3 3 3 1 2 3
+
+ return stack([rows, cols], 2)
+
+ Parameters
+ ----------
+ height, width: int
+ The height and width of feature map
+ dtype: torch.dtype
+ the data type of returned value
+ device: torch.device
+ the device of returned value
+
+ Return
+ ------
+ 2D absolute positions: torch.Tensor
+ The shape is (height, width, 2),
+ where 2 represents a 2D position (row, col).
+ '''
+ rows = torch.arange(height, dtype=dtype, device=device).view(
+ height, 1).repeat(1, width)
+ cols = torch.arange(width, dtype=dtype, device=device).view(
+ 1, width).repeat(height, 1)
+ return torch.stack([rows, cols], 2)
+
+
+class METHOD:
+ """define iRPE method IDs
+ We divide the implementation of CROSS into CROSS_ROWS and CROSS_COLS.
+
+ """
+ EUCLIDEAN = 0
+ QUANT = 1
+ PRODUCT = 3
+ CROSS = 4
+ CROSS_ROWS = 41
+ CROSS_COLS = 42
+
+
+# Define a mapping from METHOD_ID to Python function
+_METHOD_FUNC = {
+ METHOD.EUCLIDEAN: _rp_2d_euclidean,
+ METHOD.QUANT: _rp_2d_quant,
+ METHOD.PRODUCT: _rp_2d_product,
+ METHOD.CROSS_ROWS: _rp_2d_cross_rows,
+ METHOD.CROSS_COLS: _rp_2d_cross_cols,
+}
+
+
+def get_num_buckets(method, alpha, beta, gamma):
+ """ Get number of buckets storing relative position encoding.
+ The buckets does not contain `skip` token.
+
+ Parameters
+ ----------
+ method: METHOD
+ The method ID of image relative position encoding.
+ alpha, beta, gamma: float
+ The coefficients of piecewise index function.
+
+ Returns
+ -------
+ num_buckets: int
+ The number of buckets storing relative position encoding.
+ """
+ beta_int = int(beta)
+ if method == METHOD.PRODUCT:
+ # IDs in [0, (2 * beta_int + 1)^2) for Product method
+ num_buckets = (2 * beta_int + 1) ** 2
+ else:
+ # IDs in [-beta_int, beta_int] except of Product method
+ num_buckets = 2 * beta_int + 1
+ return num_buckets
+
+
+# (method, alpha, beta, gamma) -> (bucket_ids, num_buckets, height, width)
+BUCKET_IDS_BUF = dict()
+
+
+@torch.no_grad()
+def get_bucket_ids_2d_without_skip(method, height, width,
+ alpha, beta, gamma,
+ dtype=torch.long, device=torch.device('cpu')):
+ """Get bucket IDs for image relative position encodings without skip token
+
+ Parameters
+ ----------
+ method: METHOD
+ The method ID of image relative position encoding.
+ height, width: int
+ The height and width of the feature map.
+ The sequence length is equal to `height * width`.
+ alpha, beta, gamma: float
+ The coefficients of piecewise index function.
+ dtype: torch.dtype
+ the data type of returned `bucket_ids`
+ device: torch.device
+ the device of returned `bucket_ids`
+
+ Returns
+ -------
+ bucket_ids: torch.Tensor, dtype: long
+ The bucket IDs which index to corresponding encodings.
+ The shape of `bucket_ids` is (skip + L, skip + L),
+ where `L = height * wdith`.
+ num_buckets: int
+ The number of buckets including `skip` token.
+ L: int
+ The sequence length
+ """
+
+ key = (method, alpha, beta, gamma, dtype, device)
+ value = BUCKET_IDS_BUF.get(key, None)
+ if value is None or value[-2] < height or value[-1] < width:
+ if value is None:
+ max_height, max_width = height, width
+ else:
+ max_height = max(value[-2], height)
+ max_width = max(value[-1], width)
+ # relative position encoding mapping function
+ func = _METHOD_FUNC.get(method, None)
+ if func is None:
+ raise NotImplementedError(
+ f"[Error] The method ID {method} does not exist.")
+ pos = get_absolute_positions(max_height, max_width, dtype, device)
+
+ # compute the offset of a pair of 2D relative positions
+ max_L = max_height * max_width
+ pos1 = pos.view((max_L, 1, 2))
+ pos2 = pos.view((1, max_L, 2))
+ # diff: shape of (L, L, 2)
+ diff = pos1 - pos2
+
+ # bucket_ids: shape of (L, L)
+ bucket_ids = func(diff, alpha=alpha, beta=beta,
+ gamma=gamma, dtype=dtype)
+ beta_int = int(beta)
+ if method != METHOD.PRODUCT:
+ bucket_ids += beta_int
+ bucket_ids = bucket_ids.view(
+ max_height, max_width, max_height, max_width)
+
+ num_buckets = get_num_buckets(method, alpha, beta, gamma)
+ value = (bucket_ids, num_buckets, height, width)
+ BUCKET_IDS_BUF[key] = value
+ L = height * width
+ bucket_ids = value[0][:height, :width, :height, :width].reshape(L, L)
+ num_buckets = value[1]
+
+ return bucket_ids, num_buckets, L
+
+
+@torch.no_grad()
+def get_bucket_ids_2d(method, height, width,
+ skip, alpha, beta, gamma,
+ dtype=torch.long, device=torch.device('cpu')):
+ """Get bucket IDs for image relative position encodings
+
+ Parameters
+ ----------
+ method: METHOD
+ The method ID of image relative position encoding.
+ height, width: int
+ The height and width of the feature map.
+ The sequence length is equal to `height * width`.
+ skip: int
+ The number of skip token before spatial tokens.
+ When skip is 0, no classification token.
+ When skip is 1, there is a classification token before spatial tokens.
+ When skip > 1, there are `skip` extra tokens before spatial tokens.
+ alpha, beta, gamma: float
+ The coefficients of piecewise index function.
+ dtype: torch.dtype
+ the data type of returned `bucket_ids`
+ device: torch.device
+ the device of returned `bucket_ids`
+
+ Returns
+ -------
+ bucket_ids: torch.Tensor, dtype: long
+ The bucket IDs which index to corresponding encodings.
+ The shape of `bucket_ids` is (skip + L, skip + L),
+ where `L = height * wdith`.
+ num_buckets: int
+ The number of buckets including `skip` token.
+ """
+ bucket_ids, num_buckets, L = get_bucket_ids_2d_without_skip(method, height, width,
+ alpha, beta, gamma,
+ dtype, device)
+
+ # add an extra encoding (id = num_buckets) for the classification token
+ if skip > 0:
+ new_bids = bucket_ids.new_empty(size=(skip + L, skip + L))
+
+ # if extra token exists, we add extra bucket as its encoding.
+ extra_bucket_id = num_buckets
+ num_buckets += 1
+
+ new_bids[:skip] = extra_bucket_id
+ new_bids[:, :skip] = extra_bucket_id
+ new_bids[skip:, skip:] = bucket_ids
+
+ bucket_ids = new_bids
+ bucket_ids = bucket_ids.contiguous()
+ return bucket_ids, num_buckets
+
+
+class iRPE(nn.Module):
+ """The implementation of image relative position encoding (excluding Cross method).
+
+ Parameters
+ ----------
+ head_dim: int
+ The dimension for each head.
+ num_heads: int
+ The number of parallel attention heads.
+ mode: str or None
+ The mode of image relative position encoding.
+ Choices: [None, 'bias', 'contextual']
+ method: METHOD
+ The method ID of image relative position encoding.
+ The `METHOD` class is defined in `irpe.py`.
+ transposed: bool
+ Whether to transpose the input feature.
+ For iRPE on queries or keys, transposed should be `True`.
+ For iRPE on values, transposed should be `False`.
+ num_buckets: int
+ The number of buckets, which store encodings.
+ initializer: None or an inplace function
+ [Optional] The initializer to `lookup_table`.
+ Initalize `lookup_table` as zero by default.
+ rpe_config: RPEConfig
+ The config generated by the function `get_single_rpe_config`.
+ """
+ # a buffer to store bucket index
+ # (key, rp_bucket, _ctx_rp_bucket_flatten)
+ _rp_bucket_buf = (None, None, None)
+
+ def __init__(self, head_dim, num_heads=8,
+ mode=None, method=None,
+ transposed=True, num_buckets=None,
+ initializer=None, rpe_config=None):
+ super().__init__()
+ self.num_heads = num_heads
+ self.head_dim = head_dim
+
+ # relative position
+ assert mode in [None, 'bias', 'contextual']
+ self.mode = mode
+
+ assert method is not None, 'method should be a METHOD ID rather than None'
+ self.method = method
+
+ self.transposed = transposed
+ self.num_buckets = num_buckets
+
+ if initializer is None:
+ def initializer(x): return None
+ self.initializer = initializer
+
+ self.reset_parameters()
+
+ self.rpe_config = rpe_config
+
+ @torch.no_grad()
+ def reset_parameters(self):
+ # initialize the parameters of iRPE
+ if self.transposed:
+ if self.mode == 'bias':
+ self.lookup_table_bias = nn.Parameter(
+ torch.zeros(self.num_heads, self.num_buckets))
+ self.initializer(self.lookup_table_bias)
+ elif self.mode == 'contextual':
+ # shared and initialized from vit
+ pass
+ else:
+ if self.mode == 'bias':
+ raise NotImplementedError(
+ "[Error] Bias non-transposed RPE does not exist.")
+ elif self.mode == 'contextual':
+ raise ValueError('may not work, check')
+
+ def forward(self, x, height=None, width=None):
+ """forward function for iRPE.
+
+ Parameters
+ ----------
+ x: torch.Tensor
+ Input Tensor whose shape is (B, H, L, head_dim),
+ where B is batch size,
+ H is the number of heads,
+ L is the sequence length,
+ equal to height * width (+1 if class token exists)
+ head_dim is the dimension of each head
+
+ Returns
+ -------
+ rpe_encoding: torch.Tensor
+ image Relative Position Encoding,
+ whose shape is (B, H, L, L)
+ """
+ rp_bucket, self._ctx_rp_bucket_flatten = \
+ self._get_rp_bucket(x, height=height, width=width)
+
+ if self.transposed:
+ return self.forward_rpe_transpose(x, rp_bucket)
+ return self.forward_rpe_no_transpose(x, rp_bucket)
+
+ def _get_rp_bucket(self, x, height=None, width=None):
+ """Get relative position encoding buckets IDs corresponding the input shape
+
+ Parameters
+ ----------
+ x: torch.Tensor
+ Input Tensor whose shape is (B, H, L, head_dim),
+ where B is batch size,
+ H is the number of heads,
+ L is the sequence length,
+ equal to height * width (+1 if class token exists)
+ head_dim is the dimension of each head
+ height: int or None
+ [Optional] The height of the input
+ If not defined, height = floor(sqrt(L))
+ width: int or None
+ [Optional] The width of the input
+ If not defined, width = floor(sqrt(L))
+
+ Returns
+ -------
+ rp_bucket: torch.Tensor
+ relative position encoding buckets IDs
+ The shape is (L, L)
+ _ctx_rp_bucket_flatten: torch.Tensor or None
+ It is a private tensor for efficient computation.
+ """
+ B, H, L, D = x.shape
+ device = x.device
+ if height is None:
+ E = int(math.sqrt(L))
+ height = width = E
+ key = (height, width, device)
+ # use buffer if the spatial shape and device is not changable.
+
+ if self._rp_bucket_buf[0] == key:
+ return self._rp_bucket_buf[1:3]
+
+ skip = L - height * width
+ config = self.rpe_config
+ if RPEIndexFunction is not None and self.mode == 'contextual' and self.transposed:
+ # RPEIndexFunction uses int32 index.
+ dtype = torch.int32
+ else:
+ dtype = torch.long
+ rp_bucket, num_buckets = get_bucket_ids_2d(method=self.method,
+ height=height, width=width,
+ skip=skip, alpha=config.alpha,
+ beta=config.beta, gamma=config.gamma,
+ dtype=dtype, device=device)
+ assert num_buckets == self.num_buckets
+
+ # transposed contextual
+ _ctx_rp_bucket_flatten = None
+ if self.mode == 'contextual' and self.transposed:
+ if RPEIndexFunction is None:
+ offset = torch.arange(0, L * self.num_buckets, self.num_buckets,
+ dtype=rp_bucket.dtype, device=rp_bucket.device).view(-1, 1)
+ _ctx_rp_bucket_flatten = (rp_bucket + offset).flatten()
+ self._rp_bucket_buf = (key, rp_bucket, _ctx_rp_bucket_flatten)
+ return rp_bucket, _ctx_rp_bucket_flatten
+
+ def forward_rpe_transpose(self, x, rp_bucket):
+ """Forward function for iRPE (transposed version)
+ This version is utilized by RPE on Query or Key
+
+ Parameters
+ ----------
+ x: torch.Tensor
+ Input Tensor whose shape is (B, H, L, head_dim),
+ where B is batch size,
+ H is the number of heads,
+ L is the sequence length,
+ equal to height * width (+1 if class token exists)
+ head_dim is the dimension of each head
+ rp_bucket: torch.Tensor
+ relative position encoding buckets IDs
+ The shape is (L, L)
+
+ Weights
+ -------
+ lookup_table_bias: torch.Tensor
+ The shape is (H or 1, num_buckets)
+
+ or
+
+ lookup_table_weight: torch.Tensor
+ The shape is (H or 1, head_dim, num_buckets)
+
+ Returns
+ -------
+ output: torch.Tensor
+ Relative position encoding on queries or keys.
+ The shape is (B or 1, H, L, L),
+ where D is the output dimension for each head.
+ """
+
+ B = len(x) # batch_size
+ L_query, L_key = rp_bucket.shape
+ if self.mode == 'bias':
+ return self.lookup_table_bias[:, rp_bucket.flatten()]. \
+ view(1, self.num_heads, L_query, L_key)
+
+ elif self.mode == 'contextual':
+ """
+ ret[b, h, i, j] = lookup_table_weight[b, h, i, rp_bucket[i, j]]
+
+ ret[b, h, i * L_key + j] = \
+ lookup_table[b, h, i * num_buckets + rp_buckets[i, j]]
+
+ computational cost
+ ------------------
+ matmul: B * H * L_query * head_dim * num_buckets
+ index: L_query + L_query * L_key + B * H * L_query * L_key
+ total: O(B * H * L_query * (head_dim * num_buckets + L_key))
+ """
+ if RPEIndexFunction is not None:
+ return RPEIndexFunction.apply(x, rp_bucket)
+ else:
+ return x.flatten(2)[:, :, self._ctx_rp_bucket_flatten]. \
+ view(B, -1, L_query, L_key)
+
+ def forward_rpe_no_transpose(self, x, rp_bucket):
+ """Forward function for iRPE (non-transposed version)
+ This version is utilized by RPE on Value.
+
+ Parameters
+ ----------
+ x: torch.Tensor
+ Input Tensor whose shape is (B, H, L, head_dim),
+ where B is batch size,
+ H is the number of heads,
+ L is the sequence length,
+ equal to height * width (+1 if class token exists)
+ head_dim is the dimension of each head
+ rp_bucket: torch.Tensor
+ relative position encoding buckets IDs
+ The shape is (L, L)
+
+ Weights
+ -------
+ lookup_table_weight: torch.Tensor
+ The shape is (H or 1, num_buckets, head_dim)
+
+ Returns
+ -------
+ output: torch.Tensor
+ Relative position encoding on values.
+ The shape is (B, H, L, D),
+ where D is the output dimension for each head.
+ """
+
+ B = len(x) # batch_size
+ L_query, L_key = rp_bucket.shape
+ assert self.mode == 'contextual', "Only support contextual \
+version in non-transposed version"
+ weight = self.lookup_table_weight[:, rp_bucket.flatten()]. \
+ view(self.num_heads, L_query, L_key, self.head_dim)
+ # (H, L_query, B, L_key) @ (H, L_query, L_key, D) = (H, L_query, B, D)
+ # -> (B, H, L_query, D)
+ return torch.matmul(x.permute(1, 2, 0, 3), weight).permute(2, 0, 1, 3)
+
+ def __repr__(self):
+ return 'iRPE(head_dim={rpe.head_dim}, num_heads={rpe.num_heads}, \
+mode="{rpe.mode}", method={rpe.method}, transposed={rpe.transposed}, \
+num_buckets={rpe.num_buckets}, initializer={rpe.initializer}, \
+rpe_config={rpe.rpe_config})'.format(rpe=self)
+
+
+class iRPE_Cross(nn.Module):
+ """The implementation of image relative position encoding (specific for Cross method).
+
+ Parameters
+ ----------
+ head_dim: int
+ The dimension for each head.
+ num_heads: int
+ The number of parallel attention heads.
+ mode: str or None
+ The mode of image relative position encoding.
+ Choices: [None, 'bias', 'contextual']
+ method: METHOD
+ The method ID of image relative position encoding.
+ The `METHOD` class is defined in `irpe.py`.
+ transposed: bool
+ Whether to transpose the input feature.
+ For iRPE on queries or keys, transposed should be `True`.
+ For iRPE on values, transposed should be `False`.
+ num_buckets: int
+ The number of buckets, which store encodings.
+ initializer: None or an inplace function
+ [Optional] The initializer to `lookup_table`.
+ Initalize `lookup_table` as zero by default.
+ rpe_config: RPEConfig
+ The config generated by the function `get_single_rpe_config`.
+ """
+
+ def __init__(self, method, **kwargs):
+ super().__init__()
+ assert method == METHOD.CROSS
+ self.rp_rows = iRPE(**kwargs, method=METHOD.CROSS_ROWS)
+ self.rp_cols = iRPE(**kwargs, method=METHOD.CROSS_COLS)
+
+ def forward(self, x, height=None, width=None):
+ """forward function for iRPE.
+ Compute encoding on horizontal and vertical directions separately,
+ then summarize them.
+
+ Parameters
+ ----------
+ x: torch.Tensor
+ Input Tensor whose shape is (B, H, L, head_dim),
+ where B is batch size,
+ H is the number of heads,
+ L is the sequence length,
+ equal to height * width (+1 if class token exists)
+ head_dim is the dimension of each head
+ height: int or None
+ [Optional] The height of the input
+ If not defined, height = floor(sqrt(L))
+ width: int or None
+ [Optional] The width of the input
+ If not defined, width = floor(sqrt(L))
+
+ Returns
+ -------
+ rpe_encoding: torch.Tensor
+ Image Relative Position Encoding,
+ whose shape is (B, H, L, L)
+ """
+
+ rows = self.rp_rows(x, height=height, width=width)
+ cols = self.rp_cols(x, height=height, width=width)
+ return rows + cols
+
+ def __repr__(self):
+ return 'iRPE_Cross(head_dim={rpe.head_dim}, \
+num_heads={rpe.num_heads}, mode="{rpe.mode}", method={rpe.method}, \
+transposed={rpe.transposed}, num_buckets={rpe.num_buckets}, \
+initializer={rpe.initializer}, \
+rpe_config={rpe.rpe_config})'.format(rpe=self.rp_rows)
+
+
+def get_single_rpe_config(ratio=1.9,
+ method=METHOD.PRODUCT,
+ mode='contextual',
+ shared_head=True,
+ skip=0):
+ """Get the config of single relative position encoding
+
+ Parameters
+ ----------
+ ratio: float
+ The ratio to control the number of buckets.
+ method: METHOD
+ The method ID of image relative position encoding.
+ The `METHOD` class is defined in `irpe.py`.
+ mode: str or None
+ The mode of image relative position encoding.
+ Choices: [None, 'bias', 'contextual']
+ shared_head: bool
+ Whether to share weight among different heads.
+ skip: int
+ The number of skip token before spatial tokens.
+ When skip is 0, no classification token.
+ When skip is 1, there is a classification token before spatial tokens.
+ When skip > 1, there are `skip` extra tokens before spatial tokens.
+
+ Returns
+ -------
+ config: RPEConfig
+ The config of single relative position encoding.
+ """
+ config = edict()
+ # whether to share encodings across different heads
+ config.shared_head = shared_head
+ # mode: None, bias, contextual
+ config.mode = mode
+ # method: None, Bias, Quant, Cross, Product
+ config.method = method
+ # the coefficients of piecewise index function
+ config.alpha = 1 * ratio
+ config.beta = 2 * ratio
+ config.gamma = 8 * ratio
+
+ # set the number of buckets
+ config.num_buckets = get_num_buckets(method,
+ config.alpha,
+ config.beta,
+ config.gamma)
+ # add extra bucket for `skip` token (e.g. class token)
+ if skip > 0:
+ config.num_buckets += 1
+ return config
+
+
+def get_rpe_config(ratio=1.9,
+ method=METHOD.PRODUCT,
+ mode='contextual',
+ shared_head=True,
+ skip=0,
+ rpe_on='k'):
+ """Get the config of relative position encoding on queries, keys and values
+
+ Parameters
+ ----------
+ ratio: float
+ The ratio to control the number of buckets.
+ method: METHOD or str
+ The method ID (or name) of image relative position encoding.
+ The `METHOD` class is defined in `irpe.py`.
+ mode: str or None
+ The mode of image relative position encoding.
+ Choices: [None, 'bias', 'contextual']
+ shared_head: bool
+ Whether to share weight among different heads.
+ skip: int
+ The number of skip token before spatial tokens.
+ When skip is 0, no classification token.
+ When skip is 1, there is a classification token before spatial tokens.
+ When skip > 1, there are `skip` extra tokens before spatial tokens.
+ rpe_on: str
+ Where RPE attaches.
+ "q": RPE on queries
+ "k": RPE on keys
+ "v": RPE on values
+ "qk": RPE on queries and keys
+ "qkv": RPE on queries, keys and values
+
+ Returns
+ -------
+ config: RPEConfigs
+ config.rpe_q: the config of relative position encoding on queries
+ config.rpe_k: the config of relative position encoding on keys
+ config.rpe_v: the config of relative position encoding on values
+ """
+
+ # alias
+ if isinstance(method, str):
+ method_mapping = dict(
+ euc=METHOD.EUCLIDEAN,
+ quant=METHOD.QUANT,
+ cross=METHOD.CROSS,
+ product=METHOD.PRODUCT,
+ )
+ method = method_mapping[method.lower()]
+ if mode == 'ctx':
+ mode = 'contextual'
+ config = edict()
+ # relative position encoding on queries, keys and values
+ kwargs = dict(
+ ratio=ratio,
+ method=method,
+ mode=mode,
+ shared_head=shared_head,
+ skip=skip,
+ )
+ config.rpe_q = get_single_rpe_config(**kwargs) if 'q' in rpe_on else None
+ config.rpe_k = get_single_rpe_config(**kwargs) if 'k' in rpe_on else None
+ config.rpe_v = get_single_rpe_config(**kwargs) if 'v' in rpe_on else None
+ return config
+
+
+def build_rpe(config, head_dim, num_heads):
+ """Build iRPE modules on queries, keys and values.
+
+ Parameters
+ ----------
+ config: RPEConfigs
+ config.rpe_q: the config of relative position encoding on queries
+ config.rpe_k: the config of relative position encoding on keys
+ config.rpe_v: the config of relative position encoding on values
+ None when RPE is not used.
+ head_dim: int
+ The dimension for each head.
+ num_heads: int
+ The number of parallel attention heads.
+
+ Returns
+ -------
+ modules: a list of nn.Module
+ The iRPE Modules on [queries, keys, values].
+ None when RPE is not used.
+ """
+ if config is None:
+ return None, None, None
+ rpes = [config.rpe_q, config.rpe_k, config.rpe_v]
+ transposeds = [True, True, False]
+
+ def _build_single_rpe(rpe, transposed):
+ if rpe is None:
+ return None
+
+ rpe_cls = iRPE if rpe.method != METHOD.CROSS else iRPE_Cross
+ return rpe_cls(
+ head_dim=head_dim,
+ num_heads=1 if rpe.shared_head else num_heads,
+ mode=rpe.mode,
+ method=rpe.method,
+ transposed=transposed,
+ num_buckets=rpe.num_buckets,
+ rpe_config=rpe,
+ )
+ return [_build_single_rpe(rpe, transposed)
+ for rpe, transposed in zip(rpes, transposeds)]
+