fix code style bug

alibaba · Jan 2, 2025 · 67c0e39 · 67c0e39
1 parent 265b845
commit 67c0e39
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 6 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -4,7 +4,7 @@ repos:
     hooks:
       - id: insert-license
         files: \.py$
-        args: ["--license-filepath", "data/.license_header.txt", "--use-current-year"]
+        args: ["--license-filepath", "data/.license_header.txt", "--allow-past-years"]
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.7.1
     hooks:

diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2024, Alibaba Group;
+# Copyright (c) 2024-2025, Alibaba Group;
 # 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

diff --git a/scripts/pyre_check.py b/scripts/pyre_check.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2024, Alibaba Group;
+# Copyright (c) 2024-2025, Alibaba Group;
 # 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

diff --git a/setup.py b/setup.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2024, Alibaba Group;
+# Copyright (c) 2024-2025, Alibaba Group;
 # 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

diff --git a/tzrec/acc/aot_utils.py b/tzrec/acc/aot_utils.py
@@ -32,6 +32,7 @@
 
 # skip default bound check which is not allow by aot
 if "ENABLE_AOT" in os.environ:
+    # pyre-ignore [8]
     IntNBitTableBatchedEmbeddingBagsCodegen.__init__ = functools.partialmethod(
         IntNBitTableBatchedEmbeddingBagsCodegen.__init__,
         bounds_check_mode=BoundsCheckMode.NONE,
@@ -45,9 +46,10 @@
     del decomposition_table[aten._softmax.out]
 
 
+# pyre-ignore [56]
 @register_decomposition(aten._softmax)
 @out_wrapper()
-def _softmax(x: torch.Tensor, dim: int, half_to_float: bool):
+def _softmax(x: torch.Tensor, dim: int, half_to_float: bool) -> torch.Tensor:
     # eager softmax returns a contiguous tensor. Ensure that decomp also returns
     # a contiguous tensor.
     x = x.contiguous()
@@ -67,7 +69,7 @@ def _softmax(x: torch.Tensor, dim: int, half_to_float: bool):
 
 def export_model_aot(
     model: nn.Module, data: Dict[str, torch.Tensor], save_dir: str
-) -> None:
+) -> torch.export.ExportedProgram:
     """Export aot model.
 
     Args:

diff --git a/tzrec/models/model.py b/tzrec/models/model.py
@@ -241,6 +241,7 @@ def forward(
 class ScriptWrapperAOT(ScriptWrapper):
     """Model inference wrapper for aot export."""
 
+    # pyre-ignore [14]
     def forward(
         self,
         data: Dict[str, torch.Tensor],