Use itertools.pairwise when constructing MLP network.

As mentioned by
USTC-KnowledgeComputingLab/qmb#14 (comment)
we could use `itertools.pairwise` to improve the readability of codes during
constructing MLP network.

PR Tracking at: USTC-KnowledgeComputingLab/qmb#23

Author: hzhangxyz

qmb/crossmlp.py

@@ -2,6 +2,7 @@
 This file implements a cross MLP network.
 """
 
+import itertools
 import typing
 import torch
 from .bitspack import unpack_int
@@ -48,7 +49,7 @@ def __init__(self, dim_input: int, dim_output: int, hidden_size: tuple[int, ...]
         self.depth: int = len(hidden_size)
 
         dimensions: list[int] = [dim_input] + list(hidden_size) + [dim_output]
-        linears: list[torch.nn.Module] = [select_linear_layer(i, j) for i, j in zip(dimensions[:-1], dimensions[1:])]
+        linears: list[torch.nn.Module] = [select_linear_layer(i, j) for i, j in itertools.pairwise(dimensions)]
         modules: list[torch.nn.Module] = [module for linear in linears for module in (linear, torch.nn.SiLU())][:-1]
         self.model: torch.nn.Module = torch.nn.Sequential(*modules)
 

qmb/mlp.py

@@ -2,6 +2,7 @@
 This file implements the MLP network from https://arxiv.org/pdf/2109.12606 with the sampling method introduced in https://arxiv.org/pdf/2408.07625.
 """
 
+import itertools
 import torch
 from .bitspack import pack_int, unpack_int
 
@@ -47,7 +48,7 @@ def __init__(self, dim_input: int, dim_output: int, hidden_size: tuple[int, ...]
         self.depth: int = len(hidden_size)
 
         dimensions: list[int] = [dim_input] + list(hidden_size) + [dim_output]
-        linears: list[torch.nn.Module] = [select_linear_layer(i, j) for i, j in zip(dimensions[:-1], dimensions[1:])]
+        linears: list[torch.nn.Module] = [select_linear_layer(i, j) for i, j in itertools.pairwise(dimensions)]
         modules: list[torch.nn.Module] = [module for linear in linears for module in (linear, torch.nn.SiLU())][:-1]
         self.model: torch.nn.Module = torch.nn.Sequential(*modules)