onn transform pass

.gitignore

@@ -172,4 +172,6 @@ mase-trainer/
 test-trainer/
 
 # DiffLogic: tutorial files
-docs/tutorials/difflogic/data-mnist/
+docs/tutorials/difflogic/data-mnist/
+
+test/self

src/chop/actions/search/search_space/nas_bert.py

@@ -1199,9 +1199,11 @@ def forward(
         # If a 2D or 3D attention mask is provided for the cross-attention
         # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
         if self.config.is_decoder and encoder_hidden_states is not None:
-            encoder_batch_size, encoder_sequence_length, _ = (
-                encoder_hidden_states.size()
-            )
+            (
+                encoder_batch_size,
+                encoder_sequence_length,
+                _,
+            ) = encoder_hidden_states.size()
             encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
             if encoder_attention_mask is None:
                 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)

src/chop/distributed/tensor/_dispatch.py

@@ -201,8 +201,9 @@ def default_tensor(spec: _DTensorSpec) -> torch.Tensor:
                     # did not already construct one
                     random._rng_tracker = random.OffsetBasedRNGTracker(mesh.device_type)
 
-                first_arg, first_local_arg = cast(dtensor.DTensor, args[0]), cast(
-                    torch.Tensor, local_tensor_args[0]
+                first_arg, first_local_arg = (
+                    cast(dtensor.DTensor, args[0]),
+                    cast(torch.Tensor, local_tensor_args[0]),
                 )
                 rng_context = (
                     random._rng_tracker._distribute_region(first_arg._spec)

src/chop/ir/onnx/mase_onnx_graph.py

@@ -9,7 +9,6 @@
 
 
 class MaseOnnxGraph:
-
     def __init__(
         self,
         model_proto: onnx.onnx_ml_pb2.ModelProto,

src/chop/models/cnv/cnv.py

@@ -4,12 +4,8 @@
 from typing import Any
 import numpy as np
 
-from chop.nn.quantized.modules.conv2d import (
-    Conv2dBinaryResidualSign,
-)
-from chop.nn.quantized.modules.linear import (
-    LinearBinaryResidualSign,
-)
+from chop.nn.quantized.modules.conv2d import Conv2dBinaryResidualSign
+from chop.nn.quantized.modules.linear import LinearBinaryResidualSign
 from chop.models.utils import register_mase_model, register_mase_checkpoint
 
 """

src/chop/nn/backward/modules/__init__.py

@@ -1,6 +1,4 @@
-from .linear import (
-    CustomLinear,
-)
+from .linear import CustomLinear
 
 
 custom_module_map = {

src/chop/nn/optical/__init__.py

@@ -0,0 +1 @@
+from .modules import optical_module_map

src/chop/nn/optical/modules/__init__.py

@@ -0,0 +1,11 @@
+from .morr_linear import AllPassMORRCirculantLinear
+from .morr_conv2d import AllPassMORRCirculantConv2d
+
+# from ..triton_modules.morr_linear_mem import TritonMemMORRLinear
+
+
+optical_module_map = {
+    "linear_morr": AllPassMORRCirculantLinear,
+    "conv2d_morr": AllPassMORRCirculantConv2d,
+    # "linear_morr_triton": TritonMemMORRLinear,
+}

src/chop/nn/optical/modules/base_layer.py

@@ -0,0 +1,76 @@
+"""
+Description:
+Author: Jiaqi Gu (jqgu@utexas.edu)
+Date: 2021-06-08 18:55:05
+LastEditors: Jiaqi Gu (jqgu@utexas.edu)
+LastEditTime: 2021-06-08 18:55:05
+"""
+
+from typing import Any, Dict, Optional
+import torch
+from torch import nn
+from torch.types import Device
+
+__all__ = ["ONNBaseLayer"]
+
+
+class ONNBaseLayer(nn.Module):
+    def __init__(self, *args, device: Device = torch.device("cpu"), **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        # cuda or cpu, defaults to cpu
+        self.device = device
+
+    def build_parameters(self) -> None:
+        raise NotImplementedError
+
+    def reset_parameters(self) -> None:
+        raise NotImplementedError
+
+    @classmethod
+    def from_layer(cls, layer: nn.Module, *args, **kwargs) -> nn.Module:
+        raise NotImplementedError
+
+    def get_num_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    def enable_fast_forward(self) -> None:
+        self.fast_forward_flag = True
+
+    def disable_fast_forward(self) -> None:
+        self.fast_forward_flag = False
+
+    def set_phase_variation(
+        self, noise_std: float, random_state: Optional[int] = None
+    ) -> None:
+        self.phase_noise_std = noise_std
+
+    def set_gamma_noise(
+        self, noise_std: float, random_state: Optional[int] = None
+    ) -> None:
+        self.gamma_noise_std = noise_std
+
+    def set_crosstalk_factor(self, crosstalk_factor: float) -> None:
+        self.crosstalk_factor = crosstalk_factor
+
+    def set_weight_bitwidth(self, w_bit: int) -> None:
+        self.w_bit = w_bit
+
+    def set_input_bitwidth(self, in_bit: int) -> None:
+        self.in_bit = in_bit
+
+    def load_parameters(self, param_dict: Dict[str, Any]) -> None:
+        """
+        description: update parameters based on this parameter dictionary\\
+        param param_dict {dict of dict} {param_name: param_tensor, ...}
+        """
+        for name, param in param_dict.items():
+            getattr(self, name).data.copy_(param)
+
+    def switch_mode_to(self, mode: str) -> None:
+        self.mode = mode
+
+    def forward(self, x):
+        raise NotImplementedError
+
+    def extra_repr(self) -> str:
+        return ""